<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcs

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"math/rand"
	"net"
	"net/url"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/api/googleapi"
)

// A bucket that wraps another, calling its methods in a retry loop with
// randomized exponential backoff.
type retryBucket struct {
	maxSleep time.Duration
	wrapped  Bucket
}

func newRetryBucket(
		maxSleep time.Duration,
		wrapped Bucket) (b Bucket) {
	b = &retryBucket{
		maxSleep: maxSleep,
		wrapped:  wrapped,
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func shouldRetry(err error) (b bool) {
	// HTTP 50x errors.
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code >= 500 && typed.Code < 600 {
			b = true
			return
		}
	}

	// HTTP 429 errors (GCS uses these for rate limiting).
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code == 429 {
			b = true
			return
		}
	}

	// Network errors, which tend to show up transiently when doing lots of
	// operations in parallel. For example:
	//
	//     dial tcp 74.125.203.95:443: too many open files
	//
	if _, ok := err.(*net.OpError); ok {
		b = true
		return
	}

	// The HTTP package returns ErrUnexpectedEOF in several places. This seems to
	// come up when the server terminates the connection in the middle of an
	// object read.
	if err == io.ErrUnexpectedEOF {
		b = true
		return
	}

	// The HTTP library also appears to leak EOF errors from... somewhere in its
	// guts as URL errors sometimes.
	if urlErr, ok := err.(*url.Error); ok {
		if urlErr.Err == io.EOF {
			b = true
			return
		}
	}

	// Sometimes the HTTP package helpfully encapsulates the real error in a URL
	// error.
	if urlErr, ok := err.(*url.Error); ok {
		b = shouldRetry(urlErr.Err)
		return
	}

	return
}

// Choose an appropriate delay for exponential backoff, given that we have
// already slept the given number of times for this logical request.
func chooseDelay(prevSleepCount uint) (d time.Duration) {
	const baseDelay = time.Millisecond

	// Choose a a delay in [0, 2^prevSleepCount * baseDelay).
	d = (1 << prevSleepCount) * baseDelay
	d = time.Duration(rand.Int63n(int64(d)))

	return
}

// Exponential backoff for a function that might fail.
//
// This is essentially what is described in the "Best practices" section of the
// "Upload Objects" docs:
//
//     https://cloud.google.com/storage/docs/json_api/v1/how-tos/upload
//
// with the following exceptions:
//
//  *  We perform backoff for all operations.
//
//  *  The random component scales with the delay, so that the first sleep
//     cannot be as long as one second. The algorithm used matches the
//     description at http://en.wikipedia.org/wiki/Exponential_backoff.
//
//  *  We retry more types of errors; see shouldRetry above.
//
// State for total sleep time and number of previous sleeps is housed outside
// of this function to allow it to be "resumed" by multiple invocations of
// retryObjectReader.Read.
func expBackoff(
		ctx context.Context,
		desc string,
		maxSleep time.Duration,
		f func() error,
		prevSleepCount *uint,
		prevSleepDuration *time.Duration) (err error) {
	for {
		// Make an attempt. Stop if successful.
		err = f()
		if err == nil {
			return
		}

		// Do we want to retry?
		if !shouldRetry(err) {
			// Special case: don't spam up the logs for EOF, which io.Reader returns
			// in the normal course of things.
			if err != io.EOF {
				err = fmt.Errorf("not retrying %s: %w", desc, err)
			}

			return
		}

		// Choose a a delay.
		d := chooseDelay(*prevSleepCount)
		*prevSleepCount++

		// Are we out of credit?
		if *prevSleepDuration+d > maxSleep {
			// Return the most recent error.
			return
		}

		// Sleep, returning early if cancelled.
		select {
		case <-ctx.Done():
			// On cancellation, return the last error we saw.
			return

		case <-time.After(d):
			*prevSleepDuration += d
			continue
		}
	}
}

// Like expBackoff, but assumes that we've never slept before (and won't need
// to sleep again).
func oneShotExpBackoff(
		ctx context.Context,
		desc string,
		maxSleep time.Duration,
		f func() error) (err error) {
	var prevSleepCount uint
	var prevSleepDuration time.Duration

	err = expBackoff(
		ctx,
		desc,
		maxSleep,
		f,
		&prevSleepCount,
		&prevSleepDuration)

	return
}

////////////////////////////////////////////////////////////////////////
// Read support
////////////////////////////////////////////////////////////////////////

type retryObjectReader struct {
	bucket *retryBucket

	// The context we should watch when sleeping for retries.
	ctx context.Context

	// What we are trying to read.
	name       string
	generation int64
	byteRange  ByteRange

	// nil when we start or have seen a permanent error.
	wrapped io.ReadCloser

	// If we've seen an error that we shouldn't retry for, this will be non-nil
	// and should be returned permanently.
	permanentErr error

	// The number of times we've slept so far, and the total amount of time we've
	// spent sleeping.
	sleepCount    uint
	sleepDuration time.Duration
}

// Set up the wrapped reader.
func (rc *retryObjectReader) setUpWrapped() (err error) {
	// Call through to create the reader.
	req := &ReadObjectRequest{
		Name:       rc.name,
		Generation: rc.generation,
		Range:      &rc.byteRange,
	}

	wrapped, err := rc.bucket.wrapped.NewReader(rc.ctx, req)
	if err != nil {
		return
	}

	rc.wrapped = wrapped
	return
}

// Set up the wrapped reader if necessary, and make one attempt to read through
// it.
//
// Clears the wrapped reader on error.
func (rc *retryObjectReader) readOnce(p []byte) (n int, err error) {
	// Set up the wrapped reader if it's not already around.
	if rc.wrapped == nil {
		err = rc.setUpWrapped()
		if err != nil {
			return
		}
	}

	// Attempt to read from it.
	n, err = rc.wrapped.Read(p)
	if err != nil {
		rc.wrapped.Close()
		rc.wrapped = nil
		return
	}

	return
}

// Invariant: we never return an error from this function unless we've given up
// on retrying. In particular, we won't return a short read because the wrapped
// reader returned a short read and an error.
func (rc *retryObjectReader) Read(p []byte) (n int, err error) {
	// Whatever we do, accumulate the bytes that we're returning to the user.
	defer func() {
		if n < 0 {
			panic(fmt.Sprintf("Negative byte count: %d", n))
		}

		rc.byteRange.Start += uint64(n)
	}()

	// If we've already decided on a permanent error, return that.
	if rc.permanentErr != nil {
		err = rc.permanentErr
		return
	}

	// If we let an error escape below, it must be a permanent one.
	defer func() {
		if err != nil {
			rc.permanentErr = err
		}
	}()

	// We will repeatedly make single attempts until we get a successful request.
	// Don't forget to accumulate the result each time.
	tryOnce := func() (err error) {
		var bytesRead int
		bytesRead, err = rc.readOnce(p)
		n += bytesRead
		p = p[bytesRead:]

		return
	}

	err = expBackoff(
		rc.ctx,
		fmt.Sprintf("Read(%q, %d)", rc.name, rc.generation),
		rc.bucket.maxSleep,
		tryOnce,
		&rc.sleepCount,
		&rc.sleepDuration)

	return
}

func (rc *retryObjectReader) Close() (err error) {
	// If we don't have a wrapped reader, there is nothing useful that we can or
	// need to do here.
	if rc.wrapped == nil {
		return
	}

	// Call through.
	err = rc.wrapped.Close()

	return
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (rb *retryBucket) Name() (name string) {
	name = rb.wrapped.Name()
	return
}

func (rb *retryBucket) NewReader(
		ctx context.Context,
		req *ReadObjectRequest) (rc io.ReadCloser, err error) {
	// If the user specified the latest generation, we need to figure out what
	// that is so that we can create a reader that knows how to keep a stable
	// generation despite retrying repeatedly.
	var generation int64 = req.Generation
	var sleepCount uint
	var sleepDuration time.Duration

	if generation == 0 {
		findGeneration := func() (err error) {
			o, err := rb.wrapped.StatObject(
				ctx,
				&StatObjectRequest{
					Name: req.Name,
				})

			if err != nil {
				return
			}

			generation = o.Generation
			return
		}

		err = expBackoff(
			ctx,
			fmt.Sprintf("FindLatestGeneration(%q)", req.Name),
			rb.maxSleep,
			findGeneration,
			&sleepCount,
			&sleepDuration)

		if err != nil {
			return
		}
	}

	// Choose an appropriate byte range.
	byteRange := ByteRange{0, math.MaxUint64}
	if req.Range != nil {
		byteRange = *req.Range
	}

	// Now that we know what generation we're looking for, return an appropriate
	// reader that knows how to retry when the connection fails. Make sure to
	// inherit the time spent sleeping above.
	rc = &retryObjectReader{
		bucket: rb,
		ctx:    ctx,

		name:       req.Name,
		generation: generation,
		byteRange:  byteRange,

		sleepCount:    sleepCount,
		sleepDuration: sleepDuration,
	}

	return
}

func (rb *retryBucket) CreateObject(
		ctx context.Context,
		req *CreateObjectRequest) (o *Object, err error) {
	var seeker io.ReadSeeker
	if readSeeker, ok := req.Contents.(io.ReadSeeker); ok {
		seeker = readSeeker
	} else {
		// We can't simply replay the request multiple times, because the first
		// attempt might exhaust some of the req.Contents reader, leaving
		// missing contents for the second attempt.
		//
		// So, copy out all contents and create a copy of the request that we
		// will modify to serve from memory for each call.
		data, err := ioutil.ReadAll(req.Contents)
		if err != nil {
			err = fmt.Errorf("ioutil.ReadAll: %v", err)
			return nil, err
		}
		seeker = bytes.NewReader(data)
	}

	reqCopy := *req

	// Call through with that request.
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CreateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			if _, err = seeker.Seek(0, io.SeekStart); err != nil {
				return
			}
			reqCopy.Contents = seeker
			o, err = rb.wrapped.CreateObject(ctx, &reqCopy)
			return
		})

	return
}

func (rb *retryBucket) CopyObject(
		ctx context.Context,
		req *CopyObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CopyObject(%q, %q)", req.SrcName, req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.CopyObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ComposeObjects(
		ctx context.Context,
		req *ComposeObjectsRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ComposeObjects(%q)", req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.ComposeObjects(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) StatObject(
		ctx context.Context,
		req *StatObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("StatObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.StatObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ListObjects(
		ctx context.Context,
		req *ListObjectsRequest) (listing *Listing, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ListObjects(%q)", req.Prefix),
		rb.maxSleep,
		func() (err error) {
			listing, err = rb.wrapped.ListObjects(ctx, req)
			return
		})
	return
}

func (rb *retryBucket) UpdateObject(
		ctx context.Context,
		req *UpdateObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("UpdateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.UpdateObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) DeleteObject(
		ctx context.Context,
		req *DeleteObjectRequest) (err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("DeleteObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			err = rb.wrapped.DeleteObject(ctx, req)
			return
		})

	return
}
<commit_msg>Unit Test for BucketManager<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gcs

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"math"
	"math/rand"
	"net"
	"net/url"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/api/googleapi"
)

// A bucket that wraps another, calling its methods in a retry loop with
// randomized exponential backoff.
type retryBucket struct {
	maxSleep time.Duration
	wrapped  Bucket
}

func newRetryBucket(
	maxSleep time.Duration,
	wrapped Bucket) (b Bucket) {
	b = &retryBucket{
		maxSleep: maxSleep,
		wrapped:  wrapped,
	}

	return
}

////////////////////////////////////////////////////////////////////////
// Helpers
////////////////////////////////////////////////////////////////////////

func shouldRetry(err error) (b bool) {
	// HTTP 50x errors.
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code >= 500 && typed.Code < 600 {
			b = true
			return
		}
	}

	// HTTP 429 errors (GCS uses these for rate limiting).
	if typed, ok := err.(*googleapi.Error); ok {
		if typed.Code == 429 {
			b = true
			return
		}
	}

	// Network errors, which tend to show up transiently when doing lots of
	// operations in parallel. For example:
	//
	//     dial tcp 74.125.203.95:443: too many open files
	//
	if _, ok := err.(*net.OpError); ok {
		b = true
		return
	}

	// The HTTP package returns ErrUnexpectedEOF in several places. This seems to
	// come up when the server terminates the connection in the middle of an
	// object read.
	if err == io.ErrUnexpectedEOF {
		b = true
		return
	}

	// The HTTP library also appears to leak EOF errors from... somewhere in its
	// guts as URL errors sometimes.
	if urlErr, ok := err.(*url.Error); ok {
		if urlErr.Err == io.EOF {
			b = true
			return
		}
	}

	// Sometimes the HTTP package helpfully encapsulates the real error in a URL
	// error.
	if urlErr, ok := err.(*url.Error); ok {
		b = shouldRetry(urlErr.Err)
		return
	}

	return
}

// Choose an appropriate delay for exponential backoff, given that we have
// already slept the given number of times for this logical request.
func chooseDelay(prevSleepCount uint) (d time.Duration) {
	const baseDelay = time.Millisecond

	// Choose a a delay in [0, 2^prevSleepCount * baseDelay).
	d = (1 << prevSleepCount) * baseDelay
	d = time.Duration(rand.Int63n(int64(d)))

	return
}

// Exponential backoff for a function that might fail.
//
// This is essentially what is described in the "Best practices" section of the
// "Upload Objects" docs:
//
//	https://cloud.google.com/storage/docs/json_api/v1/how-tos/upload
//
// with the following exceptions:
//
//   - We perform backoff for all operations.
//
//   - The random component scales with the delay, so that the first sleep
//     cannot be as long as one second. The algorithm used matches the
//     description at http://en.wikipedia.org/wiki/Exponential_backoff.
//
//   - We retry more types of errors; see shouldRetry above.
//
// State for total sleep time and number of previous sleeps is housed outside
// of this function to allow it to be "resumed" by multiple invocations of
// retryObjectReader.Read.
func expBackoff(
	ctx context.Context,
	desc string,
	maxSleep time.Duration,
	f func() error,
	prevSleepCount *uint,
	prevSleepDuration *time.Duration) (err error) {
	for {
		// Make an attempt. Stop if successful.
		err = f()
		if err == nil {
			return
		}

		// Do we want to retry?
		if !shouldRetry(err) {
			// Special case: don't spam up the logs for EOF, which io.Reader returns
			// in the normal course of things.
			if err != io.EOF {
				err = fmt.Errorf("not retrying %s: %w", desc, err)
			}

			return
		}

		// Choose a a delay.
		d := chooseDelay(*prevSleepCount)
		*prevSleepCount++

		// Are we out of credit?
		if *prevSleepDuration+d > maxSleep {
			// Return the most recent error.
			return
		}

		// Sleep, returning early if cancelled.
		select {
		case <-ctx.Done():
			// On cancellation, return the last error we saw.
			return

		case <-time.After(d):
			*prevSleepDuration += d
			continue
		}
	}
}

// Like expBackoff, but assumes that we've never slept before (and won't need
// to sleep again).
func oneShotExpBackoff(
	ctx context.Context,
	desc string,
	maxSleep time.Duration,
	f func() error) (err error) {
	var prevSleepCount uint
	var prevSleepDuration time.Duration

	err = expBackoff(
		ctx,
		desc,
		maxSleep,
		f,
		&prevSleepCount,
		&prevSleepDuration)

	return
}

////////////////////////////////////////////////////////////////////////
// Read support
////////////////////////////////////////////////////////////////////////

type retryObjectReader struct {
	bucket *retryBucket

	// The context we should watch when sleeping for retries.
	ctx context.Context

	// What we are trying to read.
	name       string
	generation int64
	byteRange  ByteRange

	// nil when we start or have seen a permanent error.
	wrapped io.ReadCloser

	// If we've seen an error that we shouldn't retry for, this will be non-nil
	// and should be returned permanently.
	permanentErr error

	// The number of times we've slept so far, and the total amount of time we've
	// spent sleeping.
	sleepCount    uint
	sleepDuration time.Duration
}

// Set up the wrapped reader.
func (rc *retryObjectReader) setUpWrapped() (err error) {
	// Call through to create the reader.
	req := &ReadObjectRequest{
		Name:       rc.name,
		Generation: rc.generation,
		Range:      &rc.byteRange,
	}

	wrapped, err := rc.bucket.wrapped.NewReader(rc.ctx, req)
	if err != nil {
		return
	}

	rc.wrapped = wrapped
	return
}

// Set up the wrapped reader if necessary, and make one attempt to read through
// it.
//
// Clears the wrapped reader on error.
func (rc *retryObjectReader) readOnce(p []byte) (n int, err error) {
	// Set up the wrapped reader if it's not already around.
	if rc.wrapped == nil {
		err = rc.setUpWrapped()
		if err != nil {
			return
		}
	}

	// Attempt to read from it.
	n, err = rc.wrapped.Read(p)
	if err != nil {
		rc.wrapped.Close()
		rc.wrapped = nil
		return
	}

	return
}

// Invariant: we never return an error from this function unless we've given up
// on retrying. In particular, we won't return a short read because the wrapped
// reader returned a short read and an error.
func (rc *retryObjectReader) Read(p []byte) (n int, err error) {
	// Whatever we do, accumulate the bytes that we're returning to the user.
	defer func() {
		if n < 0 {
			panic(fmt.Sprintf("Negative byte count: %d", n))
		}

		rc.byteRange.Start += uint64(n)
	}()

	// If we've already decided on a permanent error, return that.
	if rc.permanentErr != nil {
		err = rc.permanentErr
		return
	}

	// If we let an error escape below, it must be a permanent one.
	defer func() {
		if err != nil {
			rc.permanentErr = err
		}
	}()

	// We will repeatedly make single attempts until we get a successful request.
	// Don't forget to accumulate the result each time.
	tryOnce := func() (err error) {
		var bytesRead int
		bytesRead, err = rc.readOnce(p)
		n += bytesRead
		p = p[bytesRead:]

		return
	}

	err = expBackoff(
		rc.ctx,
		fmt.Sprintf("Read(%q, %d)", rc.name, rc.generation),
		rc.bucket.maxSleep,
		tryOnce,
		&rc.sleepCount,
		&rc.sleepDuration)

	return
}

func (rc *retryObjectReader) Close() (err error) {
	// If we don't have a wrapped reader, there is nothing useful that we can or
	// need to do here.
	if rc.wrapped == nil {
		return
	}

	// Call through.
	err = rc.wrapped.Close()

	return
}

////////////////////////////////////////////////////////////////////////
// Public interface
////////////////////////////////////////////////////////////////////////

func (rb *retryBucket) Name() (name string) {
	name = rb.wrapped.Name()
	return
}

func (rb *retryBucket) NewReader(
	ctx context.Context,
	req *ReadObjectRequest) (rc io.ReadCloser, err error) {
	// If the user specified the latest generation, we need to figure out what
	// that is so that we can create a reader that knows how to keep a stable
	// generation despite retrying repeatedly.
	var generation int64 = req.Generation
	var sleepCount uint
	var sleepDuration time.Duration

	if generation == 0 {
		findGeneration := func() (err error) {
			o, err := rb.wrapped.StatObject(
				ctx,
				&StatObjectRequest{
					Name: req.Name,
				})

			if err != nil {
				return
			}

			generation = o.Generation
			return
		}

		err = expBackoff(
			ctx,
			fmt.Sprintf("FindLatestGeneration(%q)", req.Name),
			rb.maxSleep,
			findGeneration,
			&sleepCount,
			&sleepDuration)

		if err != nil {
			return
		}
	}

	// Choose an appropriate byte range.
	byteRange := ByteRange{0, math.MaxUint64}
	if req.Range != nil {
		byteRange = *req.Range
	}

	// Now that we know what generation we're looking for, return an appropriate
	// reader that knows how to retry when the connection fails. Make sure to
	// inherit the time spent sleeping above.
	rc = &retryObjectReader{
		bucket: rb,
		ctx:    ctx,

		name:       req.Name,
		generation: generation,
		byteRange:  byteRange,

		sleepCount:    sleepCount,
		sleepDuration: sleepDuration,
	}

	return
}

func (rb *retryBucket) CreateObject(
	ctx context.Context,
	req *CreateObjectRequest) (o *Object, err error) {
	var seeker io.ReadSeeker
	if readSeeker, ok := req.Contents.(io.ReadSeeker); ok {
		seeker = readSeeker
	} else {
		// We can't simply replay the request multiple times, because the first
		// attempt might exhaust some of the req.Contents reader, leaving
		// missing contents for the second attempt.
		//
		// So, copy out all contents and create a copy of the request that we
		// will modify to serve from memory for each call.
		data, err := ioutil.ReadAll(req.Contents)
		if err != nil {
			err = fmt.Errorf("ioutil.ReadAll: %v", err)
			return nil, err
		}
		seeker = bytes.NewReader(data)
	}

	reqCopy := *req

	// Call through with that request.
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CreateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			if _, err = seeker.Seek(0, io.SeekStart); err != nil {
				return
			}
			reqCopy.Contents = seeker
			o, err = rb.wrapped.CreateObject(ctx, &reqCopy)
			return
		})

	return
}

func (rb *retryBucket) CopyObject(
	ctx context.Context,
	req *CopyObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("CopyObject(%q, %q)", req.SrcName, req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.CopyObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ComposeObjects(
	ctx context.Context,
	req *ComposeObjectsRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ComposeObjects(%q)", req.DstName),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.ComposeObjects(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) StatObject(
	ctx context.Context,
	req *StatObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("StatObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.StatObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) ListObjects(
	ctx context.Context,
	req *ListObjectsRequest) (listing *Listing, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("ListObjects(%q)", req.Prefix),
		rb.maxSleep,
		func() (err error) {
			listing, err = rb.wrapped.ListObjects(ctx, req)
			return
		})
	return
}

func (rb *retryBucket) UpdateObject(
	ctx context.Context,
	req *UpdateObjectRequest) (o *Object, err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("UpdateObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			o, err = rb.wrapped.UpdateObject(ctx, req)
			return
		})

	return
}

func (rb *retryBucket) DeleteObject(
	ctx context.Context,
	req *DeleteObjectRequest) (err error) {
	err = oneShotExpBackoff(
		ctx,
		fmt.Sprintf("DeleteObject(%q)", req.Name),
		rb.maxSleep,
		func() (err error) {
			err = rb.wrapped.DeleteObject(ctx, req)
			return
		})

	return
}
<|endoftext|>


<commit_before>// Copyright (C) 2016-2018 vLife Systems Ltd <http://vlifesystems.com>
// Licensed under an MIT licence.  Please see LICENSE.md for details.

package experiment

import (
	"fmt"

	"github.com/lawrencewoodman/ddataset"
	"github.com/lawrencewoodman/dexpr"
	"github.com/vlifesystems/rhkit/aggregator"
	"github.com/vlifesystems/rhkit/assessment"
	"github.com/vlifesystems/rhkit/description"
	"github.com/vlifesystems/rhkit/goal"
	"github.com/vlifesystems/rhkit/rule"
	"github.com/vlifesystems/rulehunter/config"
	"github.com/vlifesystems/rulehunter/progress"
	"github.com/vlifesystems/rulehunter/quitter"
	"github.com/vlifesystems/rulehunter/report"
)

type TrainMode struct {
	dataset        ddataset.Dataset
	when           *dexpr.Expr
	ruleGeneration ruleGeneration
}

type ruleGenerationDesc struct {
	Fields            []string `yaml:"fields"`
	Arithmetic        bool     `yaml:"arithmetic"`
	CombinationLength int      `yaml:"combinationLength"`
}

type ruleGeneration struct {
	fields            []string
	arithmetic        bool
	combinationLength int
}

func (rg ruleGeneration) Fields() []string {
	return rg.fields
}

func (rg ruleGeneration) Arithmetic() bool {
	return rg.arithmetic
}

type trainModeDesc struct {
	Dataset *datasetDesc `yaml:"dataset"`
	// An expression that works out whether to run the experiment for this mode
	When           string             `yaml:"when"`
	RuleGeneration ruleGenerationDesc `yaml:"ruleGeneration"`
}

func newTrainMode(
	cfg *config.Config,
	desc *trainModeDesc,
	experimentFilename string,
	aggregators []aggregator.Spec,
	goals []*goal.Goal,
	sortOrder []assessment.SortOrder,
) (*TrainMode, error) {
	d, err := makeDataset(cfg, desc.Dataset)
	if err != nil {
		return nil, fmt.Errorf("dataset: %s", err)
	}
	when, err := makeWhenExpr(desc.When)
	if err != nil {
		return nil, InvalidWhenExprError(desc.When)
	}
	return &TrainMode{
		dataset: d,
		when:    when,
		ruleGeneration: ruleGeneration{
			fields:            desc.RuleGeneration.Fields,
			arithmetic:        desc.RuleGeneration.Arithmetic,
			combinationLength: desc.RuleGeneration.CombinationLength,
		},
	}, nil
}

func (m *TrainMode) Kind() report.ModeKind {
	return report.Train
}

func (m *TrainMode) Release() error {
	if m == nil {
		return nil
	}
	return m.dataset.Release()
}

func (m *TrainMode) Dataset() ddataset.Dataset {
	return m.dataset
}

func (m *TrainMode) NumAssessRulesStages() int {
	return 4 + m.ruleGeneration.combinationLength
}

func (m *TrainMode) Process(
	e *Experiment,
	cfg *config.Config,
	pm *progress.Monitor,
	q *quitter.Quitter,
	rules []rule.Rule,
) ([]rule.Rule, error) {
	reportProgress := func(msg string, percent float64) error {
		return pm.ReportProgress(e.File.Name(), report.Train, msg, percent)
	}
	quitReceived := func() bool {
		select {
		case <-q.C:
			return true
		default:
			return false
		}
	}
	noRules := []rule.Rule{}
	rt := newRuleTracker()

	if err := reportProgress("Describing train dataset", 0); err != nil {
		return noRules, err
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}
	desc, err := description.DescribeDataset(m.Dataset())
	if err != nil {
		return noRules, fmt.Errorf("Couldn't describe train dataset: %s", err)
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}
	rt.track(rules)
	userRules := append(rules, rule.NewTrue())
	ass, err := assessRules(e, m, 1, userRules, pm, q, cfg)
	if err != nil {
		return noRules, fmt.Errorf("Couldn't assess rules: %s", err)
	}

	assessRules := func(stage int, rules []rule.Rule) error {
		newRules := rt.track(rules)
		newAss, err :=
			assessRules(e, m, stage, newRules, pm, q, cfg)
		if err != nil {
			return fmt.Errorf("Couldn't assess rules: %s", err)
		}
		ass, err = ass.Merge(newAss)
		if err != nil {
			return fmt.Errorf("Couldn't merge assessments: %s", err)
		}
		ass.Sort(e.SortOrder)
		ass.Refine()
		ass.TruncateRuleAssessments(5000)
		return nil
	}

	if err := reportProgress("Generating rules", 0); err != nil {
		return noRules, err
	}
	generatedRules, err := rule.Generate(desc, m.ruleGeneration)
	if err != nil {
		return noRules, fmt.Errorf("Couldn't generate rules: %s", err)
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}

	if err := assessRules(2, generatedRules); err != nil {
		return noRules, err
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}

	if err := reportProgress("Tweaking rules", 0); err != nil {
		return noRules, err
	}
	tweakableRules := rule.Tweak(1, ass.Rules(), desc)

	if err := assessRules(3, tweakableRules); err != nil {
		return noRules, err
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}

	if err := reportProgress("Reduce DP of rules", 0); err != nil {
		return noRules, err
	}
	reducedDPRules := rule.ReduceDP(ass.Rules())

	if err := assessRules(4, reducedDPRules); err != nil {
		return noRules, err
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}

	ruleAssessments := []*assessment.RuleAssessment{ass.RuleAssessments[0]}
	for i := 0; i < m.ruleGeneration.combinationLength; i++ {
		if err := reportProgress("Combining rules", 0); err != nil {
			return noRules, err
		}
		combinedRules := rule.Combine(ass.Rules(), 10000)
		if err := assessRules(5+i, combinedRules); err != nil {
			return noRules, err
		}
		if quitReceived() {
			return noRules, ErrQuitReceived
		}

	outerLoop:
		// Add ruleAssessment for each combinationLength
		for _, ra := range ass.RuleAssessments {
			for _, r := range combinedRules {
				_, isTrueRule := ra.Rule.(rule.True)
				if !isTrueRule && ra.Rule.String() == r.String() {
					ruleAssessments = append(ruleAssessments, ra)
					break outerLoop
				}
			}
		}
	}

	// Add the true rule assessment if missing
	trueRuleAssessment := getTrueRuleAssessment(ruleAssessments)
	if trueRuleAssessment == nil {
		trueRuleAssessment = getTrueRuleAssessment(ass.RuleAssessments)
		if trueRuleAssessment == nil {
			panic("true rule missing from assessment")
		}
		ruleAssessments = append(ruleAssessments, trueRuleAssessment)
	}
	ass.RuleAssessments = ruleAssessments
	ass.Sort(e.SortOrder)
	ass.Refine()
	// TODO: Remove ruleAssessments that have longer combinationLength than
	// previous ruleAssessment?

	r := report.New(
		report.Train,
		e.Title,
		desc,
		ass,
		e.Aggregators,
		e.SortOrder,
		e.File.Name(),
		e.Tags,
		e.Category,
	)
	if err := r.WriteJSON(cfg); err != nil {
		return noRules, fmt.Errorf("Couldn't write JSON train report: %s", err)
	}
	return ass.Rules(), nil
}

func getTrueRuleAssessment(
	ruleAssessments []*assessment.RuleAssessment,
) *assessment.RuleAssessment {
	for _, ra := range ruleAssessments {
		if _, isTrueRule := ra.Rule.(rule.True); isTrueRule {
			return ra
		}
	}
	return nil
}
<commit_msg>Improve speed of combining rules in experiment.Process<commit_after>// Copyright (C) 2016-2018 vLife Systems Ltd <http://vlifesystems.com>
// Licensed under an MIT licence.  Please see LICENSE.md for details.

package experiment

import (
	"fmt"

	"github.com/lawrencewoodman/ddataset"
	"github.com/lawrencewoodman/dexpr"
	"github.com/vlifesystems/rhkit/aggregator"
	"github.com/vlifesystems/rhkit/assessment"
	"github.com/vlifesystems/rhkit/description"
	"github.com/vlifesystems/rhkit/goal"
	"github.com/vlifesystems/rhkit/rule"
	"github.com/vlifesystems/rulehunter/config"
	"github.com/vlifesystems/rulehunter/progress"
	"github.com/vlifesystems/rulehunter/quitter"
	"github.com/vlifesystems/rulehunter/report"
)

type TrainMode struct {
	dataset        ddataset.Dataset
	when           *dexpr.Expr
	ruleGeneration ruleGeneration
}

type ruleGenerationDesc struct {
	Fields            []string `yaml:"fields"`
	Arithmetic        bool     `yaml:"arithmetic"`
	CombinationLength int      `yaml:"combinationLength"`
}

type ruleGeneration struct {
	fields            []string
	arithmetic        bool
	combinationLength int
}

func (rg ruleGeneration) Fields() []string {
	return rg.fields
}

func (rg ruleGeneration) Arithmetic() bool {
	return rg.arithmetic
}

type trainModeDesc struct {
	Dataset *datasetDesc `yaml:"dataset"`
	// An expression that works out whether to run the experiment for this mode
	When           string             `yaml:"when"`
	RuleGeneration ruleGenerationDesc `yaml:"ruleGeneration"`
}

func newTrainMode(
	cfg *config.Config,
	desc *trainModeDesc,
	experimentFilename string,
	aggregators []aggregator.Spec,
	goals []*goal.Goal,
	sortOrder []assessment.SortOrder,
) (*TrainMode, error) {
	d, err := makeDataset(cfg, desc.Dataset)
	if err != nil {
		return nil, fmt.Errorf("dataset: %s", err)
	}
	when, err := makeWhenExpr(desc.When)
	if err != nil {
		return nil, InvalidWhenExprError(desc.When)
	}
	return &TrainMode{
		dataset: d,
		when:    when,
		ruleGeneration: ruleGeneration{
			fields:            desc.RuleGeneration.Fields,
			arithmetic:        desc.RuleGeneration.Arithmetic,
			combinationLength: desc.RuleGeneration.CombinationLength,
		},
	}, nil
}

func (m *TrainMode) Kind() report.ModeKind {
	return report.Train
}

func (m *TrainMode) Release() error {
	if m == nil {
		return nil
	}
	return m.dataset.Release()
}

func (m *TrainMode) Dataset() ddataset.Dataset {
	return m.dataset
}

func (m *TrainMode) NumAssessRulesStages() int {
	return 4 + m.ruleGeneration.combinationLength
}

func (m *TrainMode) Process(
	e *Experiment,
	cfg *config.Config,
	pm *progress.Monitor,
	q *quitter.Quitter,
	rules []rule.Rule,
) ([]rule.Rule, error) {
	reportProgress := func(msg string, percent float64) error {
		return pm.ReportProgress(e.File.Name(), report.Train, msg, percent)
	}
	quitReceived := func() bool {
		select {
		case <-q.C:
			return true
		default:
			return false
		}
	}
	noRules := []rule.Rule{}
	rt := newRuleTracker()

	if err := reportProgress("Describing train dataset", 0); err != nil {
		return noRules, err
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}
	desc, err := description.DescribeDataset(m.Dataset())
	if err != nil {
		return noRules, fmt.Errorf("Couldn't describe train dataset: %s", err)
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}
	rt.track(rules)
	userRules := append(rules, rule.NewTrue())
	ass, err := assessRules(e, m, 1, userRules, pm, q, cfg)
	if err != nil {
		return noRules, fmt.Errorf("Couldn't assess rules: %s", err)
	}

	assessRules := func(
		stage int,
		rules []rule.Rule,
	) (*assessment.Assessment, error) {
		newRules := rt.track(rules)
		newAss, err :=
			assessRules(e, m, stage, newRules, pm, q, cfg)
		if err != nil {
			return nil, fmt.Errorf("Couldn't assess rules: %s", err)
		}
		ass, err = ass.Merge(newAss)
		if err != nil {
			return nil, fmt.Errorf("Couldn't merge assessments: %s", err)
		}
		ass.Sort(e.SortOrder)
		ass.Refine()
		ass = ass.TruncateRuleAssessments(10000)
		return newAss, nil
	}

	if err := reportProgress("Generating rules", 0); err != nil {
		return noRules, err
	}
	generatedRules, err := rule.Generate(desc, m.ruleGeneration)
	if err != nil {
		return noRules, fmt.Errorf("Couldn't generate rules: %s", err)
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}

	if _, err := assessRules(2, generatedRules); err != nil {
		return noRules, err
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}

	if err := reportProgress("Tweaking rules", 0); err != nil {
		return noRules, err
	}
	tweakableRules := rule.Tweak(1, ass.Rules(), desc)

	if _, err := assessRules(3, tweakableRules); err != nil {
		return noRules, err
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}

	if err := reportProgress("Reduce DP of rules", 0); err != nil {
		return noRules, err
	}
	reducedDPRules := rule.ReduceDP(ass.Rules())

	if _, err := assessRules(4, reducedDPRules); err != nil {
		return noRules, err
	}

	if quitReceived() {
		return noRules, ErrQuitReceived
	}

	ruleAssessments := []*assessment.RuleAssessment{ass.RuleAssessments[0]}
	for i := 0; i < m.ruleGeneration.combinationLength; i++ {
		if err := reportProgress("Combining rules", 0); err != nil {
			return noRules, err
		}
		combinedRules := rule.Combine(ass.Rules(), 10000)
		combinedAss, err := assessRules(5+i, combinedRules)
		if err != nil {
			return noRules, err
		}
		if quitReceived() {
			return noRules, ErrQuitReceived
		}
		combinedAss.Sort(e.SortOrder)
		combinedAss.Refine()

		// Add ruleAssessment for each combinationLength
		for _, ra := range combinedAss.RuleAssessments {
			if _, isTrueRule := ra.Rule.(rule.True); !isTrueRule {
				ruleAssessments = append(ruleAssessments, ra)
				break
			}
		}
	}

	// Add the true rule assessment if missing
	trueRuleAssessment := getTrueRuleAssessment(ruleAssessments)
	if trueRuleAssessment == nil {
		trueRuleAssessment = getTrueRuleAssessment(ass.RuleAssessments)
		if trueRuleAssessment == nil {
			panic("true rule missing from assessment")
		}
		ruleAssessments = append(ruleAssessments, trueRuleAssessment)
	}
	ass.RuleAssessments = ruleAssessments
	ass.Sort(e.SortOrder)
	ass.Refine()
	// TODO: Remove ruleAssessments that have longer combinationLength than
	// previous ruleAssessment?

	r := report.New(
		report.Train,
		e.Title,
		desc,
		ass,
		e.Aggregators,
		e.SortOrder,
		e.File.Name(),
		e.Tags,
		e.Category,
	)
	if err := r.WriteJSON(cfg); err != nil {
		return noRules, fmt.Errorf("Couldn't write JSON train report: %s", err)
	}
	return ass.Rules(), nil
}

func getTrueRuleAssessment(
	ruleAssessments []*assessment.RuleAssessment,
) *assessment.RuleAssessment {
	for _, ra := range ruleAssessments {
		if _, isTrueRule := ra.Rule.(rule.True); isTrueRule {
			return ra
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	log "github.com/Sirupsen/logrus"
)

// Storage enumerates interactions with the storage engine, and allows us to interject in-memory
// substitutes for testing.
type Storage interface {
	Bootstrap() error

	InsertJob(SubmittedJob) (uint64, error)
}

// MongoStorage is a Storage implementation that connects to a real MongoDB cluster.
type MongoStorage struct {
	Database *mgo.Database
}

// NewMongoStorage establishes a connection to the MongoDB cluster.
func NewMongoStorage(c *Context) (*MongoStorage, error) {
	session, err := mgo.Dial(c.Settings.MongoURL)
	if err != nil {
		return nil, err
	}
	return &MongoStorage{Database: session.DB("rho")}, nil
}

func (storage *MongoStorage) jobs() *mgo.Collection {
	return storage.Database.C("jobs")
}

func (storage *MongoStorage) root() *mgo.Collection {
	return storage.Database.C("root")
}

// MongoRoot contains global metadata, counters and statistics used by various storage functions.
// Exactly one instance of MongoRoot should exist in the "root" collection.
type MongoRoot struct {
	JobID uint64 `bson:"job_id"`
}

// Bootstrap creates indices and metadata objects.
func (storage *MongoStorage) Bootstrap() error {
	initial := MongoRoot{}
	var existing MongoRoot

	info, err := storage.root().Find(bson.M{}).Apply(mgo.Change{
		Update: bson.M{"$setOnInsert": &initial},
		Upsert: true,
	}, &existing)
	if err != nil {
		return err
	}
	log.WithFields(log.Fields{
		"updated": info.Updated,
		"removed": info.Removed,
	}).Debug("MongoRoot object initialized.")

	return nil
}

// Job storage

// InsertJob appends a job to the queue and returns a newly allocated job ID.
func (storage *MongoStorage) InsertJob(job SubmittedJob) (uint64, error) {
	// Assign the job a job ID.
	var root MongoRoot
	_, err := storage.root().Find(bson.M{}).Apply(mgo.Change{
		Update:    bson.M{"$inc": bson.M{"job_id": 1}},
		ReturnNew: true,
	}, &root)
	if err != nil {
		return 0, err
	}
	job.JID = root.JobID

	if err := storage.jobs().Insert(job); err != nil {
		return 0, err
	}

	return job.JID, nil
}

// NullStorage is a useful embeddable struct that can be used to mock selected storage calls without
// needing to stub out all of the ones you don't care about.
type NullStorage struct{}

// InsertJob is a no-op.
func (storage *NullStorage) InsertJob(job SubmittedJob) (uint64, error) {
	return 0, nil
}
<commit_msg>Make sure NullStorage is a Storage.<commit_after>package main

import (
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"

	log "github.com/Sirupsen/logrus"
)

// Storage enumerates interactions with the storage engine, and allows us to interject in-memory
// substitutes for testing.
type Storage interface {
	Bootstrap() error

	InsertJob(SubmittedJob) (uint64, error)
}

// MongoStorage is a Storage implementation that connects to a real MongoDB cluster.
type MongoStorage struct {
	Database *mgo.Database
}

// NewMongoStorage establishes a connection to the MongoDB cluster.
func NewMongoStorage(c *Context) (*MongoStorage, error) {
	session, err := mgo.Dial(c.Settings.MongoURL)
	if err != nil {
		return nil, err
	}
	return &MongoStorage{Database: session.DB("rho")}, nil
}

func (storage *MongoStorage) jobs() *mgo.Collection {
	return storage.Database.C("jobs")
}

func (storage *MongoStorage) root() *mgo.Collection {
	return storage.Database.C("root")
}

// MongoRoot contains global metadata, counters and statistics used by various storage functions.
// Exactly one instance of MongoRoot should exist in the "root" collection.
type MongoRoot struct {
	JobID uint64 `bson:"job_id"`
}

// Bootstrap creates indices and metadata objects.
func (storage *MongoStorage) Bootstrap() error {
	initial := MongoRoot{}
	var existing MongoRoot

	info, err := storage.root().Find(bson.M{}).Apply(mgo.Change{
		Update: bson.M{"$setOnInsert": &initial},
		Upsert: true,
	}, &existing)
	if err != nil {
		return err
	}
	log.WithFields(log.Fields{
		"updated": info.Updated,
		"removed": info.Removed,
	}).Debug("MongoRoot object initialized.")

	return nil
}

// Job storage

// InsertJob appends a job to the queue and returns a newly allocated job ID.
func (storage *MongoStorage) InsertJob(job SubmittedJob) (uint64, error) {
	// Assign the job a job ID.
	var root MongoRoot
	_, err := storage.root().Find(bson.M{}).Apply(mgo.Change{
		Update:    bson.M{"$inc": bson.M{"job_id": 1}},
		ReturnNew: true,
	}, &root)
	if err != nil {
		return 0, err
	}
	job.JID = root.JobID

	if err := storage.jobs().Insert(job); err != nil {
		return 0, err
	}

	return job.JID, nil
}

// NullStorage is a useful embeddable struct that can be used to mock selected storage calls without
// needing to stub out all of the ones you don't care about.
type NullStorage struct{}

// Bootstrap is a no-op.
func (storage *NullStorage) Bootstrap() error {
	return nil
}

// InsertJob is a no-op.
func (storage *NullStorage) InsertJob(job SubmittedJob) (uint64, error) {
	return 0, nil
}
<|endoftext|>


<commit_before>package gaerecords

/*
	Event
	----------------------------------------------------------------------
*/

// Type that enables eventing on an object.
//   // define a struct
//   type MyObject struct {
//
//     // add an event
//     OnSomething Event
//
//   }
//
//   // create an instance of our type
//   obj := new(MyObject)
//
//   // add an event listener
//   obj.OnSomething.Do(func(e *EventContext){
//     // TODO: handle the event
//   })
//
//   // Trigger the event
//   obj.OnSomething.Trigger()
type Event struct {
	Callbacks []func(*EventContext)
}

// Adds a callback func to this event.  When Trigger() is called, the func passed
// in will get called, provided no other funcs have cancelled the event beforehand.
func (e *Event) Do(f func(*EventContext)) {
	e.Callbacks = append(e.Callbacks, f)
}

// Gets whether the event has any registered callbacks or not.
func (e *Event) HasCallbacks() bool {
	return len(e.Callbacks) > 0
}

// Triggers the event with the specified arguments.
//
// If any callbacks are registered, a new EventContext is created
// and then TriggerWithContext() is called.
//
// If no callbacks are registered, Trigger() does nothing but still
// returns a usable EventContext object.
func (e *Event) Trigger(args ...interface{}) *EventContext {

	// create a new context
	var context *EventContext = new(EventContext)
	context.Args = args
	context.Cancel = false

	if !e.HasCallbacks() {
		return context
	}

	return e.TriggerWithContext(context)

}

// Triggers the event with an existing EventContext object.
//
// All funcs that have been registered with the Do() method will
// be called.
//
// If no callbacks are registered, TriggerWithContext() does nothing.
//
// If any of the funcs sets the EventContext.Cancel property to true, no
// more callbacks will be called.
//
// Trigger() returns the EventContext that was passed through each callback which is useful
// for checking if the event chain was cancelled, or if any data was collected along the way.
//
// Usually this method is called after a Before* event that produces an EventContext object.
// This allows other events (i.e. After*) to share the same context.
func (e *Event) TriggerWithContext(context *EventContext) *EventContext {

	if !e.HasCallbacks() {
		return context
	}

	for index, c := range e.Callbacks {

		// update the index
		context.Index = index

		// call the callback
		c(context)

		// do we need to cancel?
		if context.Cancel == true {
			break
		}

	}

	return context

}

/*
	EventContext
	----------------------------------------------------------------------
*/

// Type that provides context to event callbacks.
type EventContext struct {

	// Whether the event should be cancelled or not.  If set to true inside a
	// callback func, no subsequent callbacks will be called.
	Cancel bool

	// Array holding the arguments passed to Trigger() if any.
	Args []interface{}

	// The index of this callback in the chain.  Will be 0 for first callback etc.
	Index int

	data map[string]interface{}
}

// Sets some data.
func (c *EventContext) Set(key string, value interface{}) *EventContext {

	// set the value
	c.Data()[key] = value

	// chain
	return c

}

// Gets a map[string]interface{} of the data for this context.  Will return an
// empty (but non-nil) map if no data has been provided.
func (c *EventContext) Data() map[string]interface{} {
	if c.data == nil {
		c.data = make(map[string]interface{})
	}
	return c.data
}
<commit_msg>updated event documentation<commit_after>package gaerecords

/*
	Event
	----------------------------------------------------------------------
*/

// Event enables eventing on an object.
//   // define a struct
//   type MyObject struct {
//
//     // add an event
//     OnSomething Event
//
//   }
//
//   // create an instance of our type
//   obj := new(MyObject)
//
//   // add an event listener
//   obj.OnSomething.Do(func(e *EventContext){
//     // TODO: handle the event
//   })
//
//   // Trigger the event
//   obj.OnSomething.Trigger()
type Event struct {
	Callbacks []func(*EventContext)
}

// Do adds a callback func to this event.  When Trigger() is called, the func passed
// in will get called, provided no other funcs have cancelled the event beforehand.
func (e *Event) Do(f func(*EventContext)) {
	e.Callbacks = append(e.Callbacks, f)
}

// HasCallbacks gets whether the event has any registered callbacks or not.
func (e *Event) HasCallbacks() bool {
	return len(e.Callbacks) > 0
}

// Trigger triggers the event with the specified arguments.
//
// If any callbacks are registered, a new EventContext is created
// and then TriggerWithContext() is called.
//
// If no callbacks are registered, Trigger() does nothing but still
// returns a usable EventContext object.
func (e *Event) Trigger(args ...interface{}) *EventContext {

	// create a new context
	var context *EventContext = new(EventContext)
	context.Args = args
	context.Cancel = false

	if !e.HasCallbacks() {
		return context
	}

	return e.TriggerWithContext(context)

}

// TriggerWithContext triggers the event with an existing EventContext object.
//
// All funcs that have been registered with the Do() method will
// be called.
//
// If no callbacks are registered, TriggerWithContext() does nothing.
//
// If any of the funcs sets the EventContext.Cancel property to true, no
// more callbacks will be called.
//
// Trigger() returns the EventContext that was passed through each callback which is useful
// for checking if the event chain was cancelled, or if any data was collected along the way.
//
// Usually this method is called after a Before* event that produces an EventContext object.
// This allows other events (i.e. After*) to share the same context.
func (e *Event) TriggerWithContext(context *EventContext) *EventContext {

	if !e.HasCallbacks() {
		return context
	}

	for index, c := range e.Callbacks {

		// update the index
		context.Index = index

		// call the callback
		c(context)

		// do we need to cancel?
		if context.Cancel == true {
			break
		}

	}

	return context

}

/*
	EventContext
	----------------------------------------------------------------------
*/

// EventContext provides context to event callbacks.
type EventContext struct {

	// Whether the event should be cancelled or not.  If set to true inside a
	// callback func, no subsequent callbacks will be called.
	Cancel bool

	// Array holding the arguments passed to Trigger() if any.
	Args []interface{}

	// The index of this callback in the chain.  Will be 0 for first callback etc.
	Index int

	data map[string]interface{}
}

// Set sets some data.
func (c *EventContext) Set(key string, value interface{}) *EventContext {

	// set the value
	c.Data()[key] = value

	// chain
	return c

}

// Data gets a map[string]interface{} of the data for this context.  Will return an
// empty (but non-nil) map if no data has been provided.
func (c *EventContext) Data() map[string]interface{} {
	if c.data == nil {
		c.data = make(map[string]interface{})
	}
	return c.data
}
<|endoftext|>


<commit_before>package gaeutil

import(
	"bytes"
	"encoding/gob"
	"fmt"

	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
)

/*

type Foo struct {
  S string
}

foo := Foo{S:"hello"}
err1 := gaeutil.SaveAnySingleton(ctx, "Foo_007", &foo)

foo2 := Foo{}
err2 := gaeutil.LoadAnySingleton(ctx, "Foo_007", &foo2)

 */

func singletonDSKey(c context.Context, name string) *datastore.Key {
	return datastore.NewKey(c, "Singleton", name, 0, nil)
}

type Singleton struct {
	Value []byte `datastore:",noindex"`
}

func LoadSingletonFromDatastore(c context.Context, name string) ([]byte, error) {
	s := Singleton{}
	if err := datastore.Get(c, singletonDSKey(c,name), &s); err != nil {
		return nil,err  // might be datastore.ErrNoSuchEntity
	}
	return s.Value,nil
}

func SaveSingletonToDatastore(c context.Context, name string, data []byte) error {
	if len(data) > 950000 {
		return fmt.Errorf("singleton too large (name=%s, size=%d)", name, len(data))
	}
	s := Singleton{data}
	_,err := datastore.Put(c, singletonDSKey(c,name), &s)
	return err
}

func LoadSingleton(c context.Context, name string) ([]byte, error) {
	data,err := LoadSingletonFromMemcache(c,name)
	if err != nil {
		// We don't care if it was a cache miss or a deeper error - failback to datastore either way
		data,err = LoadSingletonFromDatastore(c,name)
		if err == datastore.ErrNoSuchEntity {
			return nil,nil
		}
	}
	return data,nil
}

func SaveSingleton(c context.Context, name string, data []byte) error {
	if err := SaveSingletonToDatastore(c,name,data); err != nil {
		return err
	}

	SaveSingletonToMemcache(c,name,data)  // Don't care if this fails
	return nil
}

func DeleteSingleton(ctx context.Context, name string) error {
	DeleteSingletonFromMemcache(ctx, name) // Don't care about memcache.ErrCacheMiss
	return datastore.Delete(ctx, singletonDSKey(ctx,name))
}

func LoadAnySingleton(ctx context.Context, name string, obj interface{}) error {
	myBytes,err := LoadSingleton(ctx, name)

	if err == datastore.ErrNoSuchEntity {
		// Strictly speaking, only LoadSingletonFromDatastore should expose this miss
		return nil
	} else if err != nil {
		return err
	}

	buf := bytes.NewBuffer(myBytes)

	if err := gob.NewDecoder(buf).Decode(obj); err != nil {
		return err
	}

	return nil
}

func SaveAnySingleton(ctx context.Context, name string, obj interface{}) error {
	var buf bytes.Buffer
	if err := gob.NewEncoder(&buf).Encode(obj); err != nil {
		return err
	}

	return SaveSingleton(ctx, name, buf.Bytes())
}
<commit_msg>Don't decode empty byte arrays (else you get EOF errors)<commit_after>package gaeutil

import(
	"bytes"
	"encoding/gob"
	"fmt"

	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
)

/*

type Foo struct {
  S string
}

foo := Foo{S:"hello"}
err1 := gaeutil.SaveAnySingleton(ctx, "Foo_007", &foo)

foo2 := Foo{}
err2 := gaeutil.LoadAnySingleton(ctx, "Foo_007", &foo2)

 */

func singletonDSKey(c context.Context, name string) *datastore.Key {
	return datastore.NewKey(c, "Singleton", name, 0, nil)
}

type Singleton struct {
	Value []byte `datastore:",noindex"`
}

func LoadSingletonFromDatastore(c context.Context, name string) ([]byte, error) {
	s := Singleton{}
	if err := datastore.Get(c, singletonDSKey(c,name), &s); err != nil {
		return nil,err  // might be datastore.ErrNoSuchEntity
	}
	return s.Value,nil
}

func SaveSingletonToDatastore(c context.Context, name string, data []byte) error {
	if len(data) > 950000 {
		return fmt.Errorf("singleton too large (name=%s, size=%d)", name, len(data))
	}
	s := Singleton{data}
	_,err := datastore.Put(c, singletonDSKey(c,name), &s)
	return err
}

func LoadSingleton(c context.Context, name string) ([]byte, error) {
	data,err := LoadSingletonFromMemcache(c,name)
	if err != nil {
		// We don't care if it was a cache miss or a deeper error - failback to datastore either way
		data,err = LoadSingletonFromDatastore(c,name)

		// Why swallow this error ?
		if err == datastore.ErrNoSuchEntity {
			return nil,nil
		}
	}
	return data,nil
}

func SaveSingleton(c context.Context, name string, data []byte) error {
	if err := SaveSingletonToDatastore(c,name,data); err != nil {
		return err
	}

	SaveSingletonToMemcache(c,name,data)  // Don't care if this fails
	return nil
}

func DeleteSingleton(ctx context.Context, name string) error {
	DeleteSingletonFromMemcache(ctx, name) // Don't care about memcache.ErrCacheMiss
	return datastore.Delete(ctx, singletonDSKey(ctx,name))
}

func LoadAnySingleton(ctx context.Context, name string, obj interface{}) error {
	myBytes,err := LoadSingleton(ctx, name)

	if err == datastore.ErrNoSuchEntity {
		// Debug codepath; LoadSingleton swallows this, but if we use *FromDatastore we see it
		return nil
	} else if err != nil {
		return err
	} else if myBytes == nil {
		// This happens if the object was not found; don't try to decode it.
		return nil
	}

	buf := bytes.NewBuffer(myBytes)

	if err := gob.NewDecoder(buf).Decode(obj); err != nil {
		return err
	}

	return nil
}

func SaveAnySingleton(ctx context.Context, name string, obj interface{}) error {
	var buf bytes.Buffer
	if err := gob.NewEncoder(&buf).Encode(obj); err != nil {
		return err
	}

	return SaveSingleton(ctx, name, buf.Bytes())
}
<|endoftext|>


<commit_before>package encoding

import (
	"errors"
	"fmt"
	"io"
)

// ReadPrefix reads an 8-byte length prefixes, followed by the number of bytes
// specified in the prefix. The operation is aborted if the prefix exceeds a
// specified maximum length.
func ReadPrefix(r io.Reader, maxLen uint64) ([]byte, error) {
	prefix := make([]byte, 8)
	if n, err := r.Read(prefix); err != nil || n != len(prefix) {
		return nil, errors.New("could not read length prefix")
	}
	dataLen := DecUint64(prefix)
	if dataLen > maxLen {
		return nil, fmt.Errorf("length %d exceeds maxLen of %d", dataLen, maxLen)
	}
	// read dataLen bytes
	var data []byte
	buf := make([]byte, 1024)
	var total uint64
	for total = 0; total < dataLen; {
		n, err := r.Read(buf)
		if err != nil {
			return nil, err
		}
		data = append(data, buf[:n]...)
		total += uint64(n)
	}
	if total != dataLen {
		return nil, errors.New("length mismatch")
	}
	return data, nil
}

// ReadObject reads and decodes a length-prefixed and marshalled object.
func ReadObject(r io.Reader, obj interface{}, maxLen uint64) error {
	data, err := ReadPrefix(r, maxLen)
	if err != nil {
		return err
	}
	return Unmarshal(data, obj)
}

// WritePrefix prepends data with a 4-byte length before writing it.
func WritePrefix(w io.Writer, data []byte) (int, error) {
	return w.Write(append(EncUint64(uint64(len(data))), data...))
}

// WriteObject encodes an object and prepends it with a 4-byte length before
// writing it.
func WriteObject(w io.Writer, obj interface{}) (int, error) {
	return WritePrefix(w, Marshal(obj))
}
<commit_msg>ReadPrefix: read at most dataLen bytes<commit_after>package encoding

import (
	"errors"
	"fmt"
	"io"
)

// ReadPrefix reads an 8-byte length prefixes, followed by the number of bytes
// specified in the prefix. The operation is aborted if the prefix exceeds a
// specified maximum length.
func ReadPrefix(r io.Reader, maxLen uint64) ([]byte, error) {
	prefix := make([]byte, 8)
	if n, err := r.Read(prefix); err != nil || n != len(prefix) {
		return nil, errors.New("could not read length prefix")
	}
	dataLen := DecUint64(prefix)
	if dataLen > maxLen {
		return nil, fmt.Errorf("length %d exceeds maxLen of %d", dataLen, maxLen)
	}
	// read dataLen bytes
	data := make([]byte, dataLen)
	_, err := io.ReadFull(r, data)
	return data, err
}

// ReadObject reads and decodes a length-prefixed and marshalled object.
func ReadObject(r io.Reader, obj interface{}, maxLen uint64) error {
	data, err := ReadPrefix(r, maxLen)
	if err != nil {
		return err
	}
	return Unmarshal(data, obj)
}

// WritePrefix prepends data with a 4-byte length before writing it.
func WritePrefix(w io.Writer, data []byte) (int, error) {
	return w.Write(append(EncUint64(uint64(len(data))), data...))
}

// WriteObject encodes an object and prepends it with a 4-byte length before
// writing it.
func WriteObject(w io.Writer, obj interface{}) (int, error) {
	return WritePrefix(w, Marshal(obj))
}
<|endoftext|>


<commit_before>package engines

import (
	"bytes"
	"fmt"
	"image"
	"image/draw"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"math"
	"strconv"
	"time"

	"github.com/disintegration/imaging"
	"github.com/imdario/mergo"
	imagefile "github.com/thoas/picfit/image"
	"golang.org/x/image/bmp"
	"golang.org/x/image/tiff"
)

type GoImageEngine struct {
	DefaultFormat  string
	Format         string
	DefaultQuality int
}

type ImageTransformation func(img image.Image) *image.NRGBA

var FlipTransformations = map[string]ImageTransformation{
	"h": imaging.FlipH,
	"v": imaging.FlipV,
}

var RotateTransformations = map[int]ImageTransformation{
	90:  imaging.Rotate90,
	270: imaging.Rotate270,
	180: imaging.Rotate180,
}

type Result struct {
	Paletted *image.Paletted
	Image    *image.NRGBA
	Position int
}

type Transformation func(img image.Image, width int, height int, filter imaging.ResampleFilter) *image.NRGBA

func scalingFactor(srcWidth int, srcHeight int, destWidth int, destHeight int) float64 {
	return math.Max(float64(destWidth)/float64(srcWidth), float64(destHeight)/float64(srcHeight))
}

func ImageSize(e image.Image) (int, int) {
	return e.Bounds().Max.X, e.Bounds().Max.Y
}

func (e *GoImageEngine) Scale(img image.Image, dstWidth int, dstHeight int, upscale bool, trans Transformation) *image.NRGBA {
	width, height := ImageSize(img)

	factor := scalingFactor(width, height, dstWidth, dstHeight)

	if factor < 1 || upscale {
		return trans(img, dstWidth, dstHeight, imaging.Lanczos)
	}

	return imaging.Clone(img)
}

func (e *GoImageEngine) TransformGIF(img *imagefile.ImageFile, width int, height int, options *Options, trans Transformation) ([]byte, error) {
	g, err := gif.DecodeAll(bytes.NewReader(img.Source))

	if err != nil {
		return nil, err
	}

	length := len(g.Image)
	done := make(chan *Result)
	images := make([]*image.Paletted, length)
	processed := 0

	for i := range g.Image {
		go func(paletted *image.Paletted, width int, height int, position int, trans Transformation, options *Options) {
			done <- &Result{
				Image:    e.Scale(paletted, width, height, options.Upscale, trans),
				Position: position,
				Paletted: image.NewPaletted(image.Rect(0, 0, width, height), paletted.Palette),
			}
		}(g.Image[i], width, height, i, trans, options)
	}

	for {
		select {
		case result := <-done:
			draw.Draw(result.Paletted, image.Rect(0, 0, width, height), result.Image, image.Pt(0, 0), draw.Src)

			images[result.Position] = result.Paletted

			processed++
		case <-time.After(time.Second * 5):
			break
		}

		if processed == length {
			break
		}
	}

	close(done)

	g.Image = images

	buf := &bytes.Buffer{}

	err = gif.EncodeAll(buf, g)

	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func (e *GoImageEngine) Resize(img *imagefile.ImageFile, width int, height int, options *Options) ([]byte, error) {
	if options.Format == imaging.GIF {
		content, err := e.TransformGIF(img, width, height, options, imaging.Resize)

		if err != nil {
			return nil, err
		}

		return content, nil
	}

	image, err := e.Source(img)

	if err != nil {
		return nil, err
	}

	return e.resize(image, width, height, options)
}

func (e *GoImageEngine) resize(img image.Image, width int, height int, options *Options) ([]byte, error) {
	return e.ToBytes(e.Scale(img, width, height, options.Upscale, imaging.Resize), options.Format, options.Quality)
}

func (e *GoImageEngine) Source(img *imagefile.ImageFile) (image.Image, error) {
	return imaging.Decode(bytes.NewReader(img.Source))
}

func (e *GoImageEngine) Rotate(img *imagefile.ImageFile, deg int, options *Options) ([]byte, error) {
	image, err := e.Source(img)

	if err != nil {
		return nil, err
	}

	transform, ok := RotateTransformations[deg]

	if !ok {
		return nil, fmt.Errorf("Invalid rotate transformation degree=%d is not supported", deg)
	}

	return e.ToBytes(transform(image), options.Format, options.Quality)
}

func (e *GoImageEngine) Flip(img *imagefile.ImageFile, pos string, options *Options) ([]byte, error) {
	image, err := e.Source(img)

	if err != nil {
		return nil, err
	}

	transform, ok := FlipTransformations[pos]

	if !ok {
		return nil, fmt.Errorf("Invalid flip transformation, %s is not supported", pos)
	}

	return e.ToBytes(transform(image), options.Format, options.Quality)
}

func (e *GoImageEngine) Thumbnail(img *imagefile.ImageFile, width int, height int, options *Options) ([]byte, error) {
	if options.Format == imaging.GIF {
		content, err := e.TransformGIF(img, width, height, options, imaging.Thumbnail)

		if err != nil {
			return nil, err
		}

		return content, nil
	}

	image, err := e.Source(img)

	if err != nil {
		return nil, err
	}

	return e.thumbnail(image, width, height, options)
}

func (e *GoImageEngine) thumbnail(img image.Image, width int, height int, options *Options) ([]byte, error) {
	return e.ToBytes(e.Scale(img, width, height, options.Upscale, imaging.Thumbnail), options.Format, options.Quality)
}

func (e *GoImageEngine) Fit(img *imagefile.ImageFile, width int, height int, options *Options) ([]byte, error) {
	if options.Format == imaging.GIF {
		content, err := e.TransformGIF(img, width, height, options, imaging.Thumbnail)

		if err != nil {
			return nil, err
		}

		return content, nil
	}

	image, err := e.Source(img)

	if err != nil {
		return nil, err
	}

	return e.fit(image, width, height, options)
}

func (e *GoImageEngine) fit(img image.Image, width int, height int, options *Options) ([]byte, error) {
	return e.ToBytes(e.Scale(img, width, height, options.Upscale, imaging.Fit), options.Format, options.Quality)
}

func (e *GoImageEngine) Transform(img *imagefile.ImageFile, operation *Operation, qs map[string]string) (*imagefile.ImageFile, error) {
	params := map[string]string{
		"upscale": "1",
		"h":       "0",
		"w":       "0",
		"deg":     "90",
	}

	err := mergo.Merge(&qs, params)

	if err != nil {
		return nil, err
	}

	var quality int
	var format string

	q, ok := qs["q"]

	if ok {
		quality, err := strconv.Atoi(q)

		if err != nil {
			return nil, err
		}

		if quality > 100 {
			return nil, fmt.Errorf("Quality should be <= 100")
		}
	} else {
		quality = e.DefaultQuality
	}

	format, ok = qs["fmt"]
	filepath := img.Filepath

	if ok {
		if _, ok := ContentTypes[format]; !ok {
			return nil, fmt.Errorf("Unknown format %s", format)
		}

	}

	if format == "" && e.Format != "" {
		format = e.Format
	}

	if format == "" {
		format = img.Format()
	}

	if format == "" {
		format = e.DefaultFormat
	}

	if format != img.Format() {
		index := len(filepath) - len(img.Format())

		filepath = filepath[:index] + format
	}

	file := &imagefile.ImageFile{
		Source:   img.Source,
		Key:      img.Key,
		Headers:  img.Headers,
		Filepath: filepath,
	}

	options := &Options{
		Quality: quality,
		Format:  Formats[format],
	}

	switch operation {
	case Flip:
		pos, ok := qs["pos"]

		if !ok {
			return nil, fmt.Errorf("Parameter \"pos\" not found in query string")
		}

		content, err := e.Flip(img, pos, options)

		if err != nil {
			return nil, err
		}

		file.Processed = content

		return file, err
	case Rotate:
		deg, err := strconv.Atoi(qs["deg"])

		if err != nil {
			return nil, err
		}

		content, err := e.Rotate(img, deg, options)

		if err != nil {
			return nil, err
		}

		file.Processed = content

		return file, err
	case Thumbnail, Resize, Fit:
		var upscale bool
		var w int
		var h int

		if upscale, err = strconv.ParseBool(qs["upscale"]); err != nil {
			return nil, err
		}

		if w, err = strconv.Atoi(qs["w"]); err != nil {
			return nil, err
		}

		if h, err = strconv.Atoi(qs["h"]); err != nil {
			return nil, err
		}

		options.Upscale = upscale

		switch operation {
		case Resize:
			content, err := e.Resize(img, w, h, options)

			if err != nil {
				return nil, err
			}

			file.Processed = content

			return file, err
		case Thumbnail:
			content, err := e.Thumbnail(img, w, h, options)

			if err != nil {
				return nil, err
			}

			file.Processed = content

			return file, err
		case Fit:
			content, err := e.Fit(img, w, h, options)

			if err != nil {
				return nil, err
			}

			file.Processed = content

			return file, err
			}
		}
	}

	return nil, fmt.Errorf("Operation not found for %s", operation)
}

func (e *GoImageEngine) ToBytes(img image.Image, format imaging.Format, quality int) ([]byte, error) {
	buf := &bytes.Buffer{}

	var err error

	err = encode(buf, img, format, quality)

	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func encode(w io.Writer, img image.Image, format imaging.Format, quality int) error {
	var err error
	switch format {
	case imaging.JPEG:
		var rgba *image.RGBA
		if nrgba, ok := img.(*image.NRGBA); ok {
			if nrgba.Opaque() {
				rgba = &image.RGBA{
					Pix:    nrgba.Pix,
					Stride: nrgba.Stride,
					Rect:   nrgba.Rect,
				}
			}
		}
		if rgba != nil {
			err = jpeg.Encode(w, rgba, &jpeg.Options{Quality: quality})
		} else {
			err = jpeg.Encode(w, img, &jpeg.Options{Quality: quality})
		}

	case imaging.PNG:
		err = png.Encode(w, img)
	case imaging.GIF:
		err = gif.Encode(w, img, &gif.Options{NumColors: 256})
	case imaging.TIFF:
		err = tiff.Encode(w, img, &tiff.Options{Compression: tiff.Deflate, Predictor: true})
	case imaging.BMP:
		err = bmp.Encode(w, img)
	default:
		err = imaging.ErrUnsupportedFormat
	}
	return err
}
<commit_msg>excess }<commit_after>package engines

import (
	"bytes"
	"fmt"
	"image"
	"image/draw"
	"image/gif"
	"image/jpeg"
	"image/png"
	"io"
	"math"
	"strconv"
	"time"

	"github.com/disintegration/imaging"
	"github.com/imdario/mergo"
	imagefile "github.com/thoas/picfit/image"
	"golang.org/x/image/bmp"
	"golang.org/x/image/tiff"
)

type GoImageEngine struct {
	DefaultFormat  string
	Format         string
	DefaultQuality int
}

type ImageTransformation func(img image.Image) *image.NRGBA

var FlipTransformations = map[string]ImageTransformation{
	"h": imaging.FlipH,
	"v": imaging.FlipV,
}

var RotateTransformations = map[int]ImageTransformation{
	90:  imaging.Rotate90,
	270: imaging.Rotate270,
	180: imaging.Rotate180,
}

type Result struct {
	Paletted *image.Paletted
	Image    *image.NRGBA
	Position int
}

type Transformation func(img image.Image, width int, height int, filter imaging.ResampleFilter) *image.NRGBA

func scalingFactor(srcWidth int, srcHeight int, destWidth int, destHeight int) float64 {
	return math.Max(float64(destWidth)/float64(srcWidth), float64(destHeight)/float64(srcHeight))
}

func ImageSize(e image.Image) (int, int) {
	return e.Bounds().Max.X, e.Bounds().Max.Y
}

func (e *GoImageEngine) Scale(img image.Image, dstWidth int, dstHeight int, upscale bool, trans Transformation) *image.NRGBA {
	width, height := ImageSize(img)

	factor := scalingFactor(width, height, dstWidth, dstHeight)

	if factor < 1 || upscale {
		return trans(img, dstWidth, dstHeight, imaging.Lanczos)
	}

	return imaging.Clone(img)
}

func (e *GoImageEngine) TransformGIF(img *imagefile.ImageFile, width int, height int, options *Options, trans Transformation) ([]byte, error) {
	g, err := gif.DecodeAll(bytes.NewReader(img.Source))

	if err != nil {
		return nil, err
	}

	length := len(g.Image)
	done := make(chan *Result)
	images := make([]*image.Paletted, length)
	processed := 0

	for i := range g.Image {
		go func(paletted *image.Paletted, width int, height int, position int, trans Transformation, options *Options) {
			done <- &Result{
				Image:    e.Scale(paletted, width, height, options.Upscale, trans),
				Position: position,
				Paletted: image.NewPaletted(image.Rect(0, 0, width, height), paletted.Palette),
			}
		}(g.Image[i], width, height, i, trans, options)
	}

	for {
		select {
		case result := <-done:
			draw.Draw(result.Paletted, image.Rect(0, 0, width, height), result.Image, image.Pt(0, 0), draw.Src)

			images[result.Position] = result.Paletted

			processed++
		case <-time.After(time.Second * 5):
			break
		}

		if processed == length {
			break
		}
	}

	close(done)

	g.Image = images

	buf := &bytes.Buffer{}

	err = gif.EncodeAll(buf, g)

	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func (e *GoImageEngine) Resize(img *imagefile.ImageFile, width int, height int, options *Options) ([]byte, error) {
	if options.Format == imaging.GIF {
		content, err := e.TransformGIF(img, width, height, options, imaging.Resize)

		if err != nil {
			return nil, err
		}

		return content, nil
	}

	image, err := e.Source(img)

	if err != nil {
		return nil, err
	}

	return e.resize(image, width, height, options)
}

func (e *GoImageEngine) resize(img image.Image, width int, height int, options *Options) ([]byte, error) {
	return e.ToBytes(e.Scale(img, width, height, options.Upscale, imaging.Resize), options.Format, options.Quality)
}

func (e *GoImageEngine) Source(img *imagefile.ImageFile) (image.Image, error) {
	return imaging.Decode(bytes.NewReader(img.Source))
}

func (e *GoImageEngine) Rotate(img *imagefile.ImageFile, deg int, options *Options) ([]byte, error) {
	image, err := e.Source(img)

	if err != nil {
		return nil, err
	}

	transform, ok := RotateTransformations[deg]

	if !ok {
		return nil, fmt.Errorf("Invalid rotate transformation degree=%d is not supported", deg)
	}

	return e.ToBytes(transform(image), options.Format, options.Quality)
}

func (e *GoImageEngine) Flip(img *imagefile.ImageFile, pos string, options *Options) ([]byte, error) {
	image, err := e.Source(img)

	if err != nil {
		return nil, err
	}

	transform, ok := FlipTransformations[pos]

	if !ok {
		return nil, fmt.Errorf("Invalid flip transformation, %s is not supported", pos)
	}

	return e.ToBytes(transform(image), options.Format, options.Quality)
}

func (e *GoImageEngine) Thumbnail(img *imagefile.ImageFile, width int, height int, options *Options) ([]byte, error) {
	if options.Format == imaging.GIF {
		content, err := e.TransformGIF(img, width, height, options, imaging.Thumbnail)

		if err != nil {
			return nil, err
		}

		return content, nil
	}

	image, err := e.Source(img)

	if err != nil {
		return nil, err
	}

	return e.thumbnail(image, width, height, options)
}

func (e *GoImageEngine) thumbnail(img image.Image, width int, height int, options *Options) ([]byte, error) {
	return e.ToBytes(e.Scale(img, width, height, options.Upscale, imaging.Thumbnail), options.Format, options.Quality)
}

func (e *GoImageEngine) Fit(img *imagefile.ImageFile, width int, height int, options *Options) ([]byte, error) {
	if options.Format == imaging.GIF {
		content, err := e.TransformGIF(img, width, height, options, imaging.Thumbnail)

		if err != nil {
			return nil, err
		}

		return content, nil
	}

	image, err := e.Source(img)

	if err != nil {
		return nil, err
	}

	return e.fit(image, width, height, options)
}

func (e *GoImageEngine) fit(img image.Image, width int, height int, options *Options) ([]byte, error) {
	return e.ToBytes(e.Scale(img, width, height, options.Upscale, imaging.Fit), options.Format, options.Quality)
}

func (e *GoImageEngine) Transform(img *imagefile.ImageFile, operation *Operation, qs map[string]string) (*imagefile.ImageFile, error) {
	params := map[string]string{
		"upscale": "1",
		"h":       "0",
		"w":       "0",
		"deg":     "90",
	}

	err := mergo.Merge(&qs, params)

	if err != nil {
		return nil, err
	}

	var quality int
	var format string

	q, ok := qs["q"]

	if ok {
		quality, err := strconv.Atoi(q)

		if err != nil {
			return nil, err
		}

		if quality > 100 {
			return nil, fmt.Errorf("Quality should be <= 100")
		}
	} else {
		quality = e.DefaultQuality
	}

	format, ok = qs["fmt"]
	filepath := img.Filepath

	if ok {
		if _, ok := ContentTypes[format]; !ok {
			return nil, fmt.Errorf("Unknown format %s", format)
		}

	}

	if format == "" && e.Format != "" {
		format = e.Format
	}

	if format == "" {
		format = img.Format()
	}

	if format == "" {
		format = e.DefaultFormat
	}

	if format != img.Format() {
		index := len(filepath) - len(img.Format())

		filepath = filepath[:index] + format
	}

	file := &imagefile.ImageFile{
		Source:   img.Source,
		Key:      img.Key,
		Headers:  img.Headers,
		Filepath: filepath,
	}

	options := &Options{
		Quality: quality,
		Format:  Formats[format],
	}

	switch operation {
	case Flip:
		pos, ok := qs["pos"]

		if !ok {
			return nil, fmt.Errorf("Parameter \"pos\" not found in query string")
		}

		content, err := e.Flip(img, pos, options)

		if err != nil {
			return nil, err
		}

		file.Processed = content

		return file, err
	case Rotate:
		deg, err := strconv.Atoi(qs["deg"])

		if err != nil {
			return nil, err
		}

		content, err := e.Rotate(img, deg, options)

		if err != nil {
			return nil, err
		}

		file.Processed = content

		return file, err
	case Thumbnail, Resize, Fit:
		var upscale bool
		var w int
		var h int

		if upscale, err = strconv.ParseBool(qs["upscale"]); err != nil {
			return nil, err
		}

		if w, err = strconv.Atoi(qs["w"]); err != nil {
			return nil, err
		}

		if h, err = strconv.Atoi(qs["h"]); err != nil {
			return nil, err
		}

		options.Upscale = upscale

		switch operation {
		case Resize:
			content, err := e.Resize(img, w, h, options)

			if err != nil {
				return nil, err
			}

			file.Processed = content

			return file, err
		case Thumbnail:
			content, err := e.Thumbnail(img, w, h, options)

			if err != nil {
				return nil, err
			}

			file.Processed = content

			return file, err
		case Fit:
			content, err := e.Fit(img, w, h, options)

			if err != nil {
				return nil, err
			}

			file.Processed = content

			return file, err
		}
	}

	return nil, fmt.Errorf("Operation not found for %s", operation)
}

func (e *GoImageEngine) ToBytes(img image.Image, format imaging.Format, quality int) ([]byte, error) {
	buf := &bytes.Buffer{}

	var err error

	err = encode(buf, img, format, quality)

	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

func encode(w io.Writer, img image.Image, format imaging.Format, quality int) error {
	var err error
	switch format {
	case imaging.JPEG:
		var rgba *image.RGBA
		if nrgba, ok := img.(*image.NRGBA); ok {
			if nrgba.Opaque() {
				rgba = &image.RGBA{
					Pix:    nrgba.Pix,
					Stride: nrgba.Stride,
					Rect:   nrgba.Rect,
				}
			}
		}
		if rgba != nil {
			err = jpeg.Encode(w, rgba, &jpeg.Options{Quality: quality})
		} else {
			err = jpeg.Encode(w, img, &jpeg.Options{Quality: quality})
		}

	case imaging.PNG:
		err = png.Encode(w, img)
	case imaging.GIF:
		err = gif.Encode(w, img, &gif.Options{NumColors: 256})
	case imaging.TIFF:
		err = tiff.Encode(w, img, &tiff.Options{Compression: tiff.Deflate, Predictor: true})
	case imaging.BMP:
		err = bmp.Encode(w, img)
	default:
		err = imaging.ErrUnsupportedFormat
	}
	return err
}
<|endoftext|>