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

package ebiten

import (
	"bytes"
	"fmt"
	"go/parser"
	"go/token"
	"strings"

	"github.com/hajimehoshi/ebiten/internal/graphics"
	"github.com/hajimehoshi/ebiten/internal/mipmap"
	"github.com/hajimehoshi/ebiten/internal/shader"
	"github.com/hajimehoshi/ebiten/internal/shaderir"
)

var shaderSuffix string

func init() {
	shaderSuffix = `
var __imageDstTextureSize vec2

// imageSrcTextureSize returns the destination image's texture size in pixels.
func imageDstTextureSize() vec2 {
	return __imageDstTextureSize
}
`

	shaderSuffix += fmt.Sprintf(`
var __textureSizes [%[1]d]vec2

// imageSrcTextureSize returns the source image's texture size in pixels.
// As an image is a part of internal texture, the texture is usually bigger than the image.
// The texture's size is useful when you want to calculate pixels from texels.
func imageSrcTextureSize() vec2 {
	return __textureSizes[0]
}

// The unit is the source texture's texel.
var __textureSourceOffsets [%[2]d]vec2

// The unit is the source texture's texel.
var __textureSourceRegionOrigin vec2

// The unit is the source texture's texel.
var __textureSourceRegionSize vec2

// imageSrcTextureRegion returns the source image's region (the origin and the size) on its texture.
// The unit is the source texture's texel.
//
// As an image is a part of internal texture, the image can be located at an arbitrary position on the texture.
func imageSrcTextureRegion() (vec2, vec2) {
	return __textureSourceRegionOrigin, __textureSourceRegionSize
}
`, graphics.ShaderImageNum, graphics.ShaderImageNum-1)

	for i := 0; i < graphics.ShaderImageNum; i++ {
		pos := "pos"
		if i >= 1 {
			// Convert the position in texture0's texels to the target texture texels.
			pos = fmt.Sprintf("(pos + __textureSourceOffsets[%d]) * __textureSizes[0] / __textureSizes[%d]", i-1, i)
		}
		// __t%d is a special variable for a texture variable.
		shaderSuffix += fmt.Sprintf(`
func imageSrc%[1]dAt(pos vec2) vec4 {
	// pos is the position in texels of the source texture (= 0th image's texture).
	return texture2D(__t%[1]d, %[2]s)
}

func imageSrc%[1]dBoundsAt(pos vec2) vec4 {
	// pos is the position in texels of the source texture (= 0th image's texture).
	return texture2D(__t%[1]d, %[2]s) *
		step(__textureSourceRegionOrigin.x, pos.x) *
		(1 - step(__textureSourceRegionOrigin.x + __textureSourceRegionSize.x, pos.x)) *
		step(__textureSourceRegionOrigin.y, pos.y) *
		(1 - step(__textureSourceRegionOrigin.y + __textureSourceRegionSize.y, pos.y))
}
`, i, pos)
	}

	shaderSuffix += `
func __vertex(position vec2, texCoord vec2, color vec4) (vec4, vec2, vec4) {
	return mat4(
		2/imageDstTextureSize().x, 0, 0, 0,
		0, 2/imageDstTextureSize().y, 0, 0,
		0, 0, 1, 0,
		-1, -1, 0, 1,
	) * vec4(position, 0, 1), texCoord, color
}
`
}

type Shader struct {
	shader       *mipmap.Shader
	uniformNames []string
	uniformTypes []shaderir.Type
}

func NewShader(src []byte) (*Shader, error) {
	var buf bytes.Buffer
	buf.Write(src)
	buf.WriteString(shaderSuffix)

	fs := token.NewFileSet()
	f, err := parser.ParseFile(fs, "", buf.Bytes(), parser.AllErrors)
	if err != nil {
		return nil, err
	}

	const (
		vert = "__vertex"
		frag = "Fragment"
	)
	s, err := shader.Compile(fs, f, vert, frag, graphics.ShaderImageNum)
	if err != nil {
		return nil, err
	}

	if s.VertexFunc.Block == nil {
		return nil, fmt.Errorf("ebiten: vertex shader entry point '%s' is missing", vert)
	}
	if s.FragmentFunc.Block == nil {
		return nil, fmt.Errorf("ebiten: fragment shader entry point '%s' is missing", frag)
	}

	return &Shader{
		shader:       mipmap.NewShader(s),
		uniformNames: s.UniformNames,
		uniformTypes: s.Uniforms,
	}, nil
}

func (s *Shader) Dispose() {
	s.shader.MarkDisposed()
	s.shader = nil
}

func (s *Shader) convertUniforms(uniforms map[string]interface{}) []interface{} {
	type index struct {
		resultIndex        int
		shaderUniformIndex int
	}

	names := map[string]index{}
	var idx int
	for i, n := range s.uniformNames {
		if strings.HasPrefix(n, "__") {
			continue
		}
		names[n] = index{
			resultIndex:        idx,
			shaderUniformIndex: i,
		}
		idx++
	}

	us := make([]interface{}, len(names))
	for name, idx := range names {
		if v, ok := uniforms[name]; ok {
			// TODO: Check the uniform variable types?
			us[idx.resultIndex] = v
			continue
		}

		t := s.uniformTypes[idx.shaderUniformIndex]
		v := zeroUniformValue(t)
		if v == nil {
			panic(fmt.Sprintf("ebiten: unexpected uniform variable type: %s", t.String()))
		}
		us[idx.resultIndex] = v
	}

	// TODO: Panic if uniforms include an invalid name

	return us
}

func zeroUniformValue(t shaderir.Type) interface{} {
	switch t.Main {
	case shaderir.Bool:
		return false
	case shaderir.Int:
		return 0
	case shaderir.Float:
		return float32(0)
	case shaderir.Array:
		switch t.Sub[0].Main {
		case shaderir.Bool:
			return make([]bool, t.Length)
		case shaderir.Int:
			return make([]int, t.Length)
		default:
			return make([]float32, t.FloatNum())
		}
	default:
		return make([]float32, t.FloatNum())
	}
}
<commit_msg>ebiten: Refactoring shader.go<commit_after>// Copyright 2020 The Ebiten Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package ebiten

import (
	"bytes"
	"fmt"
	"go/parser"
	"go/token"
	"strings"

	"github.com/hajimehoshi/ebiten/internal/graphics"
	"github.com/hajimehoshi/ebiten/internal/mipmap"
	"github.com/hajimehoshi/ebiten/internal/shader"
	"github.com/hajimehoshi/ebiten/internal/shaderir"
)

var shaderSuffix string

func init() {
	shaderSuffix = `
var __imageDstTextureSize vec2

// imageSrcTextureSize returns the destination image's texture size in pixels.
func imageDstTextureSize() vec2 {
	return __imageDstTextureSize
}
`

	shaderSuffix += fmt.Sprintf(`
var __textureSizes [%[1]d]vec2

// imageSrcTextureSize returns the source image's texture size in pixels.
// As an image is a part of internal texture, the texture is usually bigger than the image.
// The texture's size is useful when you want to calculate pixels from texels.
func imageSrcTextureSize() vec2 {
	return __textureSizes[0]
}

// The unit is the source texture's texel.
var __textureSourceOffsets [%[2]d]vec2

// The unit is the source texture's texel.
var __textureSourceRegionOrigin vec2

// The unit is the source texture's texel.
var __textureSourceRegionSize vec2

// imageSrcTextureRegion returns the source image's region (the origin and the size) on its texture.
// The unit is the source texture's texel.
//
// As an image is a part of internal texture, the image can be located at an arbitrary position on the texture.
func imageSrcTextureRegion() (vec2, vec2) {
	return __textureSourceRegionOrigin, __textureSourceRegionSize
}
`, graphics.ShaderImageNum, graphics.ShaderImageNum-1)

	for i := 0; i < graphics.ShaderImageNum; i++ {
		pos := "pos"
		if i >= 1 {
			// Convert the position in texture0's texels to the target texture texels.
			pos = fmt.Sprintf("(pos + __textureSourceOffsets[%d]) * __textureSizes[0] / __textureSizes[%d]", i-1, i)
		}
		// __t%d is a special variable for a texture variable.
		shaderSuffix += fmt.Sprintf(`
func imageSrc%[1]dAt(pos vec2) vec4 {
	// pos is the position in texels of the source texture (= 0th image's texture).
	return texture2D(__t%[1]d, %[2]s)
}

func imageSrc%[1]dBoundsAt(pos vec2) vec4 {
	// pos is the position in texels of the source texture (= 0th image's texture).
	return texture2D(__t%[1]d, %[2]s) *
		step(__textureSourceRegionOrigin.x, pos.x) *
		(1 - step(__textureSourceRegionOrigin.x + __textureSourceRegionSize.x, pos.x)) *
		step(__textureSourceRegionOrigin.y, pos.y) *
		(1 - step(__textureSourceRegionOrigin.y + __textureSourceRegionSize.y, pos.y))
}
`, i, pos)
	}

	shaderSuffix += `
func __vertex(position vec2, texCoord vec2, color vec4) (vec4, vec2, vec4) {
	return mat4(
		2/__imageDstTextureSize.x, 0, 0, 0,
		0, 2/__imageDstTextureSize.y, 0, 0,
		0, 0, 1, 0,
		-1, -1, 0, 1,
	) * vec4(position, 0, 1), texCoord, color
}
`
}

type Shader struct {
	shader       *mipmap.Shader
	uniformNames []string
	uniformTypes []shaderir.Type
}

func NewShader(src []byte) (*Shader, error) {
	var buf bytes.Buffer
	buf.Write(src)
	buf.WriteString(shaderSuffix)

	fs := token.NewFileSet()
	f, err := parser.ParseFile(fs, "", buf.Bytes(), parser.AllErrors)
	if err != nil {
		return nil, err
	}

	const (
		vert = "__vertex"
		frag = "Fragment"
	)
	s, err := shader.Compile(fs, f, vert, frag, graphics.ShaderImageNum)
	if err != nil {
		return nil, err
	}

	if s.VertexFunc.Block == nil {
		return nil, fmt.Errorf("ebiten: vertex shader entry point '%s' is missing", vert)
	}
	if s.FragmentFunc.Block == nil {
		return nil, fmt.Errorf("ebiten: fragment shader entry point '%s' is missing", frag)
	}

	return &Shader{
		shader:       mipmap.NewShader(s),
		uniformNames: s.UniformNames,
		uniformTypes: s.Uniforms,
	}, nil
}

func (s *Shader) Dispose() {
	s.shader.MarkDisposed()
	s.shader = nil
}

func (s *Shader) convertUniforms(uniforms map[string]interface{}) []interface{} {
	type index struct {
		resultIndex        int
		shaderUniformIndex int
	}

	names := map[string]index{}
	var idx int
	for i, n := range s.uniformNames {
		if strings.HasPrefix(n, "__") {
			continue
		}
		names[n] = index{
			resultIndex:        idx,
			shaderUniformIndex: i,
		}
		idx++
	}

	us := make([]interface{}, len(names))
	for name, idx := range names {
		if v, ok := uniforms[name]; ok {
			// TODO: Check the uniform variable types?
			us[idx.resultIndex] = v
			continue
		}

		t := s.uniformTypes[idx.shaderUniformIndex]
		v := zeroUniformValue(t)
		if v == nil {
			panic(fmt.Sprintf("ebiten: unexpected uniform variable type: %s", t.String()))
		}
		us[idx.resultIndex] = v
	}

	// TODO: Panic if uniforms include an invalid name

	return us
}

func zeroUniformValue(t shaderir.Type) interface{} {
	switch t.Main {
	case shaderir.Bool:
		return false
	case shaderir.Int:
		return 0
	case shaderir.Float:
		return float32(0)
	case shaderir.Array:
		switch t.Sub[0].Main {
		case shaderir.Bool:
			return make([]bool, t.Length)
		case shaderir.Int:
			return make([]int, t.Length)
		default:
			return make([]float32, t.FloatNum())
		}
	default:
		return make([]float32, t.FloatNum())
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 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 zoekt

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"time"

	"golang.org/x/net/context"

	"github.com/fsnotify/fsnotify"
	"github.com/google/zoekt/query"
)

type searchShard struct {
	Searcher
	mtime time.Time
}

type shardWatcher struct {
	dir string

	// Limit the number of parallel queries. Since searching is
	// CPU bound, we can't do better than #CPU queries in
	// parallel.  If we do so, we just create more memory
	// pressure.
	throttle chan struct{}

	shards map[string]*searchShard
	quit   chan struct{}
}

func loadShard(fn string) (*searchShard, error) {
	f, err := os.Open(fn)
	if err != nil {
		return nil, err
	}
	fi, err := f.Stat()
	if err != nil {
		return nil, err
	}

	iFile, err := NewIndexFile(f)
	if err != nil {
		return nil, err
	}
	s, err := NewSearcher(iFile)
	if err != nil {
		return nil, fmt.Errorf("NewSearcher(%s): %v", fn, err)
	}

	return &searchShard{
		mtime:    fi.ModTime(),
		Searcher: s,
	}, nil
}

func (s *shardWatcher) String() string {
	return fmt.Sprintf("shardWatcher(%s)", s.dir)
}

func (s *shardWatcher) scan() error {
	fs, err := filepath.Glob(filepath.Join(s.dir, "*.zoekt"))
	if err != nil {
		return err
	}

	if len(fs) == 0 {
		return fmt.Errorf("directory %s is empty", s.dir)
	}

	ts := map[string]time.Time{}
	for _, fn := range fs {
		key := filepath.Base(fn)
		fi, err := os.Lstat(fn)
		if err != nil {
			continue
		}

		ts[key] = fi.ModTime()
	}

	s.lock()
	var toLoad []string
	for k, mtime := range ts {
		if s.shards[k] == nil || s.shards[k].mtime != mtime {
			toLoad = append(toLoad, k)
		}
	}

	var toDrop []string
	// Unload deleted shards.
	for k := range s.shards {
		if _, ok := ts[k]; !ok {
			toDrop = append(toDrop, k)
		}
	}
	s.unlock()

	for _, t := range toDrop {
		log.Printf("unloading: %s", t)
		s.replace(t, nil)
	}

	for _, t := range toLoad {
		shard, err := loadShard(filepath.Join(s.dir, t))
		log.Printf("reloading: %s, err %v ", t, err)
		if err != nil {
			continue
		}
		s.replace(t, shard)
	}

	return nil
}

func (s *shardWatcher) rlock() {
	s.throttle <- struct{}{}
}

// getShards returns the currently loaded shards. The shards must be
// accessed under a rlock call.
func (s *shardWatcher) getShards() []Searcher {
	var res []Searcher
	for _, sh := range s.shards {
		res = append(res, sh)
	}
	return res
}

func (s *shardWatcher) runlock() {
	<-s.throttle
}

func (s *shardWatcher) lock() {
	n := cap(s.throttle)
	for n > 0 {
		s.throttle <- struct{}{}
		n--
	}
}

func (s *shardWatcher) unlock() {
	n := cap(s.throttle)
	for n > 0 {
		<-s.throttle
		n--
	}
}

func (s *shardWatcher) replace(key string, shard *searchShard) {
	s.lock()
	defer s.unlock()
	old := s.shards[key]
	if old != nil {
		old.Close()
	}
	if shard != nil {
		s.shards[key] = shard
	} else {
		delete(s.shards, key)
	}
}

func (s *shardWatcher) watch() error {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return err
	}
	if err := watcher.Add(s.dir); err != nil {
		return err
	}

	go func() {
		for {
			select {
			case <-watcher.Events:
				s.scan()
			case err := <-watcher.Errors:
				if err != nil {
					log.Println("watcher error:", err)
				}
			case <-s.quit:
				watcher.Close()
				return
			}
		}
	}()
	return nil
}

// NewShardedSearcher returns a searcher instance that loads all
// shards corresponding to a glob into memory.
func NewShardedSearcher(dir string) (Searcher, error) {
	ss := shardWatcher{
		dir:      dir,
		shards:   make(map[string]*searchShard),
		quit:     make(chan struct{}, 1),
		throttle: make(chan struct{}, runtime.NumCPU()),
	}

	if err := ss.scan(); err != nil {
		return nil, err
	}

	if err := ss.watch(); err != nil {
		return nil, err
	}

	return &shardedSearcher{&ss}, nil
}

// Close closes references to open files. It may be called only once.
func (ss *shardWatcher) Close() {
	close(ss.quit)
	ss.lock()
	defer ss.unlock()
	for _, s := range ss.shards {
		s.Close()
	}
}

type shardLoader interface {
	Close()
	getShards() []Searcher
	rlock()
	runlock()
	String() string
}

type shardedSearcher struct {
	shardLoader
}

func (ss *shardedSearcher) Stats() (*RepoStats, error) {
	var r RepoStats
	ss.rlock()
	defer ss.runlock()
	for _, s := range ss.shardLoader.getShards() {
		s, err := s.Stats()
		if err != nil {
			return nil, err
		}
		r.Add(s)
	}
	return &r, nil
}

func (ss *shardedSearcher) Search(ctx context.Context, pat query.Q, opts *SearchOptions) (*SearchResult, error) {
	start := time.Now()
	type res struct {
		sr  *SearchResult
		err error
	}

	aggregate := SearchResult{
		RepoURLs:      map[string]string{},
		LineFragments: map[string]string{},
	}

	// This critical section is large, but we don't want to deal with
	// searches on shards that have just been closed.
	ss.shardLoader.rlock()
	aggregate.Wait = time.Now().Sub(start)
	start = time.Now()

	// TODO - allow for canceling the query.
	shards := ss.getShards()
	all := make(chan res, len(shards))

	var childCtx context.Context
	var cancel context.CancelFunc
	if opts.MaxWallTime == 0 {
		childCtx, cancel = context.WithCancel(ctx)
	} else {
		childCtx, cancel = context.WithTimeout(ctx, opts.MaxWallTime)
	}

	defer cancel()

	for _, s := range shards {
		go func(s Searcher) {
			defer func() {
				if r := recover(); r != nil {
					log.Printf("crashed shard: %s: %s", s.String(), r)

					var r SearchResult
					r.Stats.Crashes = 1
					all <- res{&r, nil}
				}
			}()
			ms, err := s.Search(childCtx, pat, opts)
			all <- res{ms, err}
		}(s)
	}

	for range shards {
		r := <-all
		if r.err != nil {
			return nil, r.err
		}
		aggregate.Files = append(aggregate.Files, r.sr.Files...)
		aggregate.Stats.Add(r.sr.Stats)

		if len(r.sr.Files) > 0 {
			for k, v := range r.sr.RepoURLs {
				aggregate.RepoURLs[k] = v
			}
			for k, v := range r.sr.LineFragments {
				aggregate.LineFragments[k] = v
			}
		}

		if cancel != nil && aggregate.Stats.MatchCount > opts.TotalMaxMatchCount {
			cancel()
			cancel = nil
		}
	}
	ss.runlock()

	sortFilesByScore(aggregate.Files)
	aggregate.Duration = time.Now().Sub(start)
	return &aggregate, nil
}

func (ss *shardedSearcher) List(ctx context.Context, r query.Q) (*RepoList, error) {
	type res struct {
		rl  *RepoList
		err error
	}

	ss.rlock()
	shards := ss.getShards()
	shardCount := len(shards)
	all := make(chan res, shardCount)
	for _, s := range shards {
		go func(s Searcher) {
			defer func() {
				if r := recover(); r != nil {
					all <- res{
						&RepoList{Crashes: 1}, nil,
					}
				}
			}()
			ms, err := s.List(ctx, r)
			all <- res{ms, err}
		}(s)
	}

	crashes := 0
	uniq := map[string]*Repository{}
	for i := 0; i < shardCount; i++ {
		r := <-all
		if r.err != nil {
			return nil, r.err
		}
		crashes += r.rl.Crashes
		for _, r := range r.rl.Repos {
			uniq[r.Name] = r
		}
	}
	ss.runlock()

	var names []string
	for n := range uniq {
		names = append(names, n)
	}
	sort.Strings(names)

	var aggregate []*Repository
	for _, k := range names {
		aggregate = append(aggregate, uniq[k])
	}
	return &RepoList{
		Repos:   aggregate,
		Crashes: crashes,
	}, nil
}
<commit_msg>Print stacktrace for crashed shards.<commit_after>// Copyright 2016 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 zoekt

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"runtime/debug"
	"sort"
	"time"

	"golang.org/x/net/context"

	"github.com/fsnotify/fsnotify"
	"github.com/google/zoekt/query"
)

type searchShard struct {
	Searcher
	mtime time.Time
}

type shardWatcher struct {
	dir string

	// Limit the number of parallel queries. Since searching is
	// CPU bound, we can't do better than #CPU queries in
	// parallel.  If we do so, we just create more memory
	// pressure.
	throttle chan struct{}

	shards map[string]*searchShard
	quit   chan struct{}
}

func loadShard(fn string) (*searchShard, error) {
	f, err := os.Open(fn)
	if err != nil {
		return nil, err
	}
	fi, err := f.Stat()
	if err != nil {
		return nil, err
	}

	iFile, err := NewIndexFile(f)
	if err != nil {
		return nil, err
	}
	s, err := NewSearcher(iFile)
	if err != nil {
		return nil, fmt.Errorf("NewSearcher(%s): %v", fn, err)
	}

	return &searchShard{
		mtime:    fi.ModTime(),
		Searcher: s,
	}, nil
}

func (s *shardWatcher) String() string {
	return fmt.Sprintf("shardWatcher(%s)", s.dir)
}

func (s *shardWatcher) scan() error {
	fs, err := filepath.Glob(filepath.Join(s.dir, "*.zoekt"))
	if err != nil {
		return err
	}

	if len(fs) == 0 {
		return fmt.Errorf("directory %s is empty", s.dir)
	}

	ts := map[string]time.Time{}
	for _, fn := range fs {
		key := filepath.Base(fn)
		fi, err := os.Lstat(fn)
		if err != nil {
			continue
		}

		ts[key] = fi.ModTime()
	}

	s.lock()
	var toLoad []string
	for k, mtime := range ts {
		if s.shards[k] == nil || s.shards[k].mtime != mtime {
			toLoad = append(toLoad, k)
		}
	}

	var toDrop []string
	// Unload deleted shards.
	for k := range s.shards {
		if _, ok := ts[k]; !ok {
			toDrop = append(toDrop, k)
		}
	}
	s.unlock()

	for _, t := range toDrop {
		log.Printf("unloading: %s", t)
		s.replace(t, nil)
	}

	for _, t := range toLoad {
		shard, err := loadShard(filepath.Join(s.dir, t))
		log.Printf("reloading: %s, err %v ", t, err)
		if err != nil {
			continue
		}
		s.replace(t, shard)
	}

	return nil
}

func (s *shardWatcher) rlock() {
	s.throttle <- struct{}{}
}

// getShards returns the currently loaded shards. The shards must be
// accessed under a rlock call.
func (s *shardWatcher) getShards() []Searcher {
	var res []Searcher
	for _, sh := range s.shards {
		res = append(res, sh)
	}
	return res
}

func (s *shardWatcher) runlock() {
	<-s.throttle
}

func (s *shardWatcher) lock() {
	n := cap(s.throttle)
	for n > 0 {
		s.throttle <- struct{}{}
		n--
	}
}

func (s *shardWatcher) unlock() {
	n := cap(s.throttle)
	for n > 0 {
		<-s.throttle
		n--
	}
}

func (s *shardWatcher) replace(key string, shard *searchShard) {
	s.lock()
	defer s.unlock()
	old := s.shards[key]
	if old != nil {
		old.Close()
	}
	if shard != nil {
		s.shards[key] = shard
	} else {
		delete(s.shards, key)
	}
}

func (s *shardWatcher) watch() error {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		return err
	}
	if err := watcher.Add(s.dir); err != nil {
		return err
	}

	go func() {
		for {
			select {
			case <-watcher.Events:
				s.scan()
			case err := <-watcher.Errors:
				if err != nil {
					log.Println("watcher error:", err)
				}
			case <-s.quit:
				watcher.Close()
				return
			}
		}
	}()
	return nil
}

// NewShardedSearcher returns a searcher instance that loads all
// shards corresponding to a glob into memory.
func NewShardedSearcher(dir string) (Searcher, error) {
	ss := shardWatcher{
		dir:      dir,
		shards:   make(map[string]*searchShard),
		quit:     make(chan struct{}, 1),
		throttle: make(chan struct{}, runtime.NumCPU()),
	}

	if err := ss.scan(); err != nil {
		return nil, err
	}

	if err := ss.watch(); err != nil {
		return nil, err
	}

	return &shardedSearcher{&ss}, nil
}

// Close closes references to open files. It may be called only once.
func (ss *shardWatcher) Close() {
	close(ss.quit)
	ss.lock()
	defer ss.unlock()
	for _, s := range ss.shards {
		s.Close()
	}
}

type shardLoader interface {
	Close()
	getShards() []Searcher
	rlock()
	runlock()
	String() string
}

type shardedSearcher struct {
	shardLoader
}

func (ss *shardedSearcher) Stats() (*RepoStats, error) {
	var r RepoStats
	ss.rlock()
	defer ss.runlock()
	for _, s := range ss.shardLoader.getShards() {
		s, err := s.Stats()
		if err != nil {
			return nil, err
		}
		r.Add(s)
	}
	return &r, nil
}

func (ss *shardedSearcher) Search(ctx context.Context, pat query.Q, opts *SearchOptions) (*SearchResult, error) {
	start := time.Now()
	type res struct {
		sr  *SearchResult
		err error
	}

	aggregate := SearchResult{
		RepoURLs:      map[string]string{},
		LineFragments: map[string]string{},
	}

	// This critical section is large, but we don't want to deal with
	// searches on shards that have just been closed.
	ss.shardLoader.rlock()
	aggregate.Wait = time.Now().Sub(start)
	start = time.Now()

	// TODO - allow for canceling the query.
	shards := ss.getShards()
	all := make(chan res, len(shards))

	var childCtx context.Context
	var cancel context.CancelFunc
	if opts.MaxWallTime == 0 {
		childCtx, cancel = context.WithCancel(ctx)
	} else {
		childCtx, cancel = context.WithTimeout(ctx, opts.MaxWallTime)
	}

	defer cancel()

	for _, s := range shards {
		go func(s Searcher) {
			defer func() {
				if r := recover(); r != nil {
					log.Printf("crashed shard: %s: %s, %s", s.String(), r, debug.Stack())

					var r SearchResult
					r.Stats.Crashes = 1
					all <- res{&r, nil}
				}
			}()
			ms, err := s.Search(childCtx, pat, opts)
			all <- res{ms, err}
		}(s)
	}

	for range shards {
		r := <-all
		if r.err != nil {
			return nil, r.err
		}
		aggregate.Files = append(aggregate.Files, r.sr.Files...)
		aggregate.Stats.Add(r.sr.Stats)

		if len(r.sr.Files) > 0 {
			for k, v := range r.sr.RepoURLs {
				aggregate.RepoURLs[k] = v
			}
			for k, v := range r.sr.LineFragments {
				aggregate.LineFragments[k] = v
			}
		}

		if cancel != nil && aggregate.Stats.MatchCount > opts.TotalMaxMatchCount {
			cancel()
			cancel = nil
		}
	}
	ss.runlock()

	sortFilesByScore(aggregate.Files)
	aggregate.Duration = time.Now().Sub(start)
	return &aggregate, nil
}

func (ss *shardedSearcher) List(ctx context.Context, r query.Q) (*RepoList, error) {
	type res struct {
		rl  *RepoList
		err error
	}

	ss.rlock()
	shards := ss.getShards()
	shardCount := len(shards)
	all := make(chan res, shardCount)
	for _, s := range shards {
		go func(s Searcher) {
			defer func() {
				if r := recover(); r != nil {
					all <- res{
						&RepoList{Crashes: 1}, nil,
					}
				}
			}()
			ms, err := s.List(ctx, r)
			all <- res{ms, err}
		}(s)
	}

	crashes := 0
	uniq := map[string]*Repository{}
	for i := 0; i < shardCount; i++ {
		r := <-all
		if r.err != nil {
			return nil, r.err
		}
		crashes += r.rl.Crashes
		for _, r := range r.rl.Repos {
			uniq[r.Name] = r
		}
	}
	ss.runlock()

	var names []string
	for n := range uniq {
		names = append(names, n)
	}
	sort.Strings(names)

	var aggregate []*Repository
	for _, k := range names {
		aggregate = append(aggregate, uniq[k])
	}
	return &RepoList{
		Repos:   aggregate,
		Crashes: crashes,
	}, nil
}
<|endoftext|>


<commit_before>package anaconda

import (
	"net/url"
)

type Cursor struct {
	Previous_cursor     int64
	Previous_cursor_str string

	Ids []int64

	Next_cursor     int64
	Next_cursor_str string
}

type TwitterUserCursor struct {
	Previous_cursor     int64
	Previous_cursor_str string
	Next_cursor         int64
	Next_cursor_str     string
	Users               []TwitterUser
}

type Friendship struct {
	Name        string
	Id_str      string
	Id          int64
	Connections []string
	Screen_name string
}

type FollowersPage struct {
	Followers []TwitterUser
	Error     error
}

//GetFriendshipsNoRetweets s a collection of user_ids that the currently authenticated user does not want to receive retweets from.
//It does not currently support the stringify_ids parameter
func (a TwitterApi) GetFriendshipsNoRetweets() (ids []int64, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"https://api.twitter.com/1.1/friendships/no_retweets/ids.json", nil, &ids, _GET, response_ch}
	return ids, (<-response_ch).err
}

func (a TwitterApi) GetFollowersIds(v url.Values) (c Cursor, err error) {
	err = a.apiGet("https://api.twitter.com/1.1/followers/ids.json", v, &c)
	return
}

func (a TwitterApi) GetFriendsIds(v url.Values) (c Cursor, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"https://api.twitter.com/1.1/friends/ids.json", v, &c, _GET, response_ch}
	return c, (<-response_ch).err
}

func (a TwitterApi) GetFriendshipsLookup(v url.Values) (friendships []Friendship, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"http://api.twitter.com/1.1/friendships/lookup.json", v, &friendships, _GET, response_ch}
	return friendships, (<-response_ch).err
}

func (a TwitterApi) GetFriendshipsIncoming(v url.Values) (c Cursor, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"https://api.twitter.com/1.1/friendships/incoming.json", v, &c, _GET, response_ch}
	return c, (<-response_ch).err
}

func (a TwitterApi) GetFriendshipsOutgoing(v url.Values) (c Cursor, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"http://api.twitter.com/1.1/friendships/outgoing.json", v, &c, _GET, response_ch}
	return c, (<-response_ch).err
}

func (a TwitterApi) GetFollowersList(v url.Values) (c TwitterUserCursor, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"https://api.twitter.com/1.1/followers/list.json", v, &c, _GET, response_ch}
	return c, (<-response_ch).err
}

// Like GetFollowersList, but returns a channel instead of a cursor and pre-fetches the remaining results
// This channel is closed once all values have been fetched
func (a TwitterApi) GetFollowersListAll(v url.Values) (result chan FollowersPage) {

	result = make(chan FollowersPage)

	if v == nil {
		v = url.Values{}
	}
	go func(a TwitterApi, v url.Values, result chan FollowersPage) {
		// Cursor defaults to the first page ("-1")
		next_cursor := "-1"
		for {
			v.Set("cursor", next_cursor)
			c, err := a.GetFollowersList(v)

			// throttledQuery() handles all rate-limiting errors
			// if GetFollowersList() returns an error, it must be a different kind of error

			result <- FollowersPage{c.Users, err}

			next_cursor = c.Next_cursor_str
			if next_cursor == "0" {
				close(result)
				break
			}
		}
	}(a, v, result)
	return result
}
<commit_msg> Implement GetFriendsIdsAll (which returns a channel)<commit_after>package anaconda

import (
	"net/url"
)

type Cursor struct {
	Previous_cursor     int64
	Previous_cursor_str string

	Ids []int64

	Next_cursor     int64
	Next_cursor_str string
}

type TwitterUserCursor struct {
	Previous_cursor     int64
	Previous_cursor_str string
	Next_cursor         int64
	Next_cursor_str     string
	Users               []TwitterUser
}

type Friendship struct {
	Name        string
	Id_str      string
	Id          int64
	Connections []string
	Screen_name string
}

type FollowersPage struct {
	Followers []TwitterUser
	Error     error
}

//GetFriendshipsNoRetweets s a collection of user_ids that the currently authenticated user does not want to receive retweets from.
//It does not currently support the stringify_ids parameter
func (a TwitterApi) GetFriendshipsNoRetweets() (ids []int64, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"https://api.twitter.com/1.1/friendships/no_retweets/ids.json", nil, &ids, _GET, response_ch}
	return ids, (<-response_ch).err
}

func (a TwitterApi) GetFollowersIds(v url.Values) (c Cursor, err error) {
	err = a.apiGet("https://api.twitter.com/1.1/followers/ids.json", v, &c)
	return
}

func (a TwitterApi) GetFriendsIds(v url.Values) (c Cursor, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"https://api.twitter.com/1.1/friends/ids.json", v, &c, _GET, response_ch}
	return c, (<-response_ch).err
}

func (a TwitterApi) GetFriendshipsLookup(v url.Values) (friendships []Friendship, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"http://api.twitter.com/1.1/friendships/lookup.json", v, &friendships, _GET, response_ch}
	return friendships, (<-response_ch).err
}

func (a TwitterApi) GetFriendshipsIncoming(v url.Values) (c Cursor, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"https://api.twitter.com/1.1/friendships/incoming.json", v, &c, _GET, response_ch}
	return c, (<-response_ch).err
}

func (a TwitterApi) GetFriendshipsOutgoing(v url.Values) (c Cursor, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"http://api.twitter.com/1.1/friendships/outgoing.json", v, &c, _GET, response_ch}
	return c, (<-response_ch).err
}

func (a TwitterApi) GetFollowersList(v url.Values) (c TwitterUserCursor, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"https://api.twitter.com/1.1/followers/list.json", v, &c, _GET, response_ch}
	return c, (<-response_ch).err
}

// Like GetFollowersList, but returns a channel instead of a cursor and pre-fetches the remaining results
// This channel is closed once all values have been fetched
func (a TwitterApi) GetFollowersListAll(v url.Values) (result chan FollowersPage) {

	result = make(chan FollowersPage)

	if v == nil {
		v = url.Values{}
	}
	go func(a TwitterApi, v url.Values, result chan FollowersPage) {
		// Cursor defaults to the first page ("-1")
		next_cursor := "-1"
		for {
			v.Set("cursor", next_cursor)
			c, err := a.GetFollowersList(v)

			// throttledQuery() handles all rate-limiting errors
			// if GetFollowersList() returns an error, it must be a different kind of error

			result <- FollowersPage{c.Users, err}

			next_cursor = c.Next_cursor_str
			if next_cursor == "0" {
				close(result)
				break
			}
		}
	}(a, v, result)
	return result
}

// Like GetFriendsIds, but returns a channel instead of a cursor and pre-fetches the remaining results
// This channel is closed once all values have been fetched
func (a TwitterApi) GetFriendsIdsAll(v url.Values) (c Cursor, err error) {
	response_ch := make(chan response)
	a.queryQueue <- query{"https://api.twitter.com/1.1/friends/ids.json", v, &c, _GET, response_ch}
	return c, (<-response_ch).err
}
<|endoftext|>


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

package vt

import (
	"bufio"
	"compress/bzip2"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"
)

// FeedType ...
type FeedType string

const (
	// FileFeed is the feed type passed to NewFeed() for getting a feed with
	// all the files being scanned by VirusTotal.
	FileFeed FeedType = "files"
)

// A Feed represents a stream of objects received from VirusTotal via the
// feed API v3. This API allows you to get information about objects as they are
// processed by VirusTotal in real-time. Objects are sent on channel C.
type Feed struct {
	C        chan *Object
	client   *Client
	feedType FeedType
	// t is the time of the current package and n is index of the current item
	// within the package, the feed cursor is determined by the t and n.
	t                        time.Time
	n                        int64
	stop                     chan bool
	stopped                  bool
	err                      error
	missingPackagesTolerancy int
}

// FeedOption represents an option passed to a NewFeed.
type FeedOption func(*Feed) error

// FeedBufferSize specifies the size of the Feed's buffer.
func FeedBufferSize(size int) FeedOption {
	return func(f *Feed) error {
		f.C = make(chan *Object, size)
		return nil
	}
}

// FeedCursor specifies the point in time where the feed starts. Files processed
// by VirusTotal after that time will be retrieved. The cursor is a string with
// the format YYYYMMDDhhmm, indicating the date and time with minute precision.
// If a empty string is passed as cursor the current time will be used.
func FeedCursor(cursor string) FeedOption {
	return func(f *Feed) error {
		var err error
		// An empty cursor is acceptable, it's equivalent to passing no cursor
		// at all.
		if cursor == "" {
			return nil
		}
		// Cursor can be either YYYYMMDDhhmm or YYYYMMDDhhmm-N where N
		// indicates a line number within package YYYYMMDDhhmm.
		s := strings.Split(cursor, "-")
		if len(s) > 1 {
			f.n, err = strconv.ParseInt(s[1], 10, 32)
		}
		if err == nil {
			f.t, err = time.Parse("200601021504", s[0])
		}
		return err
	}
}

// NewFeed creates a Feed that receives objects from the specified type. Objects
// are send on channel C. The feed can be stopped at any moment by calling Stop.
// This example illustrates how a Feed is typically used:
//
//  feed, err := vt.Client(<api key>).NewFeed(vt.FileFeed)
//  if err != nil {
//     ... handle error
//  }
//  for fileObj := range feed.C {
//     ... do something with file object
//  }
//  if feed.Error() != nil {
//     ... feed as been stopped by some error.
//  }
//
func (cli *Client) NewFeed(t FeedType, options ...FeedOption) (*Feed, error) {
	feed := &Feed{
		client:   cli,
		feedType: t,
		t:        time.Now().UTC().Add(-30 * time.Minute),
		stop:     make(chan bool, 1),
		missingPackagesTolerancy: 1,
	}

	for _, opt := range options {
		if err := opt(feed); err != nil {
			return nil, err
		}
	}

	// If the channel hasn't been created yet with a custom buffer size by
	// WithBufferSize, let's create it with a default size.
	if feed.C == nil {
		feed.C = make(chan *Object, 1000)
	}

	go feed.retrieve()

	return feed, nil
}

// Cursor returns a string that can be passed to FeedCursor for creating a
// feed that resumes where a previous one left.
func (f *Feed) Cursor() string {
	return fmt.Sprintf("%s-%d", f.t.Format("200601021504"), f.n)
}

// Error returns any error occurred so far.
func (f *Feed) Error() error {
	return f.err
}

// Stop causes the feed to stop sending objects to the channel C. After Stop is
// called the feed still sends all the objects that it has buffered.
func (f *Feed) Stop() error {
	if !f.stopped {
		f.stopped = true
		f.stop <- true
	}
	return nil
}

// Send the object to the feed's channel, except if it was stopped.
func (f *Feed) sendToChannel(object *Object) int {
	select {
	case <-f.stop:
		return stop
	case f.C <- object:
		return ok
	}
}

// Wait for the given amount of time, but exits earlier if the feed is stopped
// during the waiting period.
func (f *Feed) wait(d time.Duration) int {
	select {
	case <-f.stop:
		return stop
	case <-time.After(d):
		return ok
	}
}

var errNoAvailableYet = errors.New("not available yet")
var errNotFound = errors.New("not found")

func (f *Feed) getObjects(packageTime string) ([]*Object, error) {

	u := URL("feeds/%s/%s", f.feedType, packageTime)

	httpResp, err := f.client.sendRequest("GET", u, nil, nil)
	if err != nil {
		return nil, err
	}
	defer httpResp.Body.Close()

	switch httpResp.StatusCode {
	case http.StatusBadRequest:
		if resp, err := f.client.parseResponse(httpResp); err != nil {
			if resp.Error.Code == "NotAvailableYet" {
				return nil, errNoAvailableYet
			}
		}
	case http.StatusNotFound:
		return nil, errNotFound
	}

	if httpResp.StatusCode != http.StatusOK {
		return nil, errors.New(httpResp.Status)
	}

	sc := bufio.NewScanner(bzip2.NewReader(httpResp.Body))
	// By default bufio.Scanner uses a buffer that is limited to a maximum size
	// defined by bufio.MaxScanBufferSize (64KB). This is too small for
	// accommodating the large JSONs stored in the feed files. So we create an
	// initial 1MB buffer and let it grow up to 10MB.
	buffer := make([]byte, 1*1024*1024)
	sc.Buffer(buffer, 10*1024*1024)

	objects := make([]*Object, 0)
	for sc.Scan() {
		obj := &Object{}
		if err := json.Unmarshal(sc.Bytes(), obj); err != nil {
			return objects, err
		}
		objects = append(objects, obj)
	}

	return objects, sc.Err()
}

func (f *Feed) retrieve() {
	waitDuration := 20 * time.Second
	missingPackages := 0
loop:
	for {
		packageTime := f.t.Format("200601021504") // YYYYMMDDhhmm
		objects, err := f.getObjects(packageTime)
		objects = objects[f.n:]
		switch err {
		case nil:
			for _, object := range objects {
				if f.sendToChannel(object) == stop {
					break loop
				}
				f.n++
			}
			f.t = f.t.Add(60 * time.Second)
			f.n = 0
			waitDuration = 20 * time.Second
			missingPackages = 0
		case errNoAvailableYet:
			// Feed package is not available yet, let's wait for 1 minute and
			// try again. If Close() is called during the waiting period it
			// exits early and breaks the loop.
			if f.wait(waitDuration) == stop {
				break loop
			}
			waitDuration *= 2
		case errNotFound:
			// The feed tolerates some missing packages, if the number of missing
			// packages is greater than missingPackagesTolerancy an error is
			// returned, if not, it tries to get the next package.
			missingPackages++
			if missingPackages > f.missingPackagesTolerancy {
				f.err = err
				break loop
			}
			f.t = f.t.Add(60 * time.Second)
		default:
			f.err = err
			break loop
		}
	}
	f.stopped = true
	close(f.C)
	close(f.stop)
}
<commit_msg>Use a default cursor of 1 hour ago, not 30 minutes (#8)<commit_after>// Copyright © 2019 The vt-go authors. All Rights Reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package vt

import (
	"bufio"
	"compress/bzip2"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"
)

// FeedType ...
type FeedType string

const (
	// FileFeed is the feed type passed to NewFeed() for getting a feed with
	// all the files being scanned by VirusTotal.
	FileFeed FeedType = "files"
)

// A Feed represents a stream of objects received from VirusTotal via the
// feed API v3. This API allows you to get information about objects as they are
// processed by VirusTotal in real-time. Objects are sent on channel C.
type Feed struct {
	C        chan *Object
	client   *Client
	feedType FeedType
	// t is the time of the current package and n is index of the current item
	// within the package, the feed cursor is determined by the t and n.
	t                        time.Time
	n                        int64
	stop                     chan bool
	stopped                  bool
	err                      error
	missingPackagesTolerancy int
}

// FeedOption represents an option passed to a NewFeed.
type FeedOption func(*Feed) error

// FeedBufferSize specifies the size of the Feed's buffer.
func FeedBufferSize(size int) FeedOption {
	return func(f *Feed) error {
		f.C = make(chan *Object, size)
		return nil
	}
}

// FeedCursor specifies the point in time where the feed starts. Files processed
// by VirusTotal after that time will be retrieved. The cursor is a string with
// the format YYYYMMDDhhmm, indicating the date and time with minute precision.
// If a empty string is passed as cursor the current time will be used.
func FeedCursor(cursor string) FeedOption {
	return func(f *Feed) error {
		var err error
		// An empty cursor is acceptable, it's equivalent to passing no cursor
		// at all.
		if cursor == "" {
			return nil
		}
		// Cursor can be either YYYYMMDDhhmm or YYYYMMDDhhmm-N where N
		// indicates a line number within package YYYYMMDDhhmm.
		s := strings.Split(cursor, "-")
		if len(s) > 1 {
			f.n, err = strconv.ParseInt(s[1], 10, 32)
		}
		if err == nil {
			f.t, err = time.Parse("200601021504", s[0])
		}
		return err
	}
}

// NewFeed creates a Feed that receives objects from the specified type. Objects
// are send on channel C. The feed can be stopped at any moment by calling Stop.
// This example illustrates how a Feed is typically used:
//
//  feed, err := vt.Client(<api key>).NewFeed(vt.FileFeed)
//  if err != nil {
//     ... handle error
//  }
//  for fileObj := range feed.C {
//     ... do something with file object
//  }
//  if feed.Error() != nil {
//     ... feed as been stopped by some error.
//  }
//
func (cli *Client) NewFeed(t FeedType, options ...FeedOption) (*Feed, error) {
	feed := &Feed{
		client:   cli,
		feedType: t,
		t:        time.Now().UTC().Add(-1 * time.Hour),
		stop:     make(chan bool, 1),
		missingPackagesTolerancy: 1,
	}

	for _, opt := range options {
		if err := opt(feed); err != nil {
			return nil, err
		}
	}

	// If the channel hasn't been created yet with a custom buffer size by
	// WithBufferSize, let's create it with a default size.
	if feed.C == nil {
		feed.C = make(chan *Object, 1000)
	}

	go feed.retrieve()

	return feed, nil
}

// Cursor returns a string that can be passed to FeedCursor for creating a
// feed that resumes where a previous one left.
func (f *Feed) Cursor() string {
	return fmt.Sprintf("%s-%d", f.t.Format("200601021504"), f.n)
}

// Error returns any error occurred so far.
func (f *Feed) Error() error {
	return f.err
}

// Stop causes the feed to stop sending objects to the channel C. After Stop is
// called the feed still sends all the objects that it has buffered.
func (f *Feed) Stop() error {
	if !f.stopped {
		f.stopped = true
		f.stop <- true
	}
	return nil
}

// Send the object to the feed's channel, except if it was stopped.
func (f *Feed) sendToChannel(object *Object) int {
	select {
	case <-f.stop:
		return stop
	case f.C <- object:
		return ok
	}
}

// Wait for the given amount of time, but exits earlier if the feed is stopped
// during the waiting period.
func (f *Feed) wait(d time.Duration) int {
	select {
	case <-f.stop:
		return stop
	case <-time.After(d):
		return ok
	}
}

var errNoAvailableYet = errors.New("not available yet")
var errNotFound = errors.New("not found")

func (f *Feed) getObjects(packageTime string) ([]*Object, error) {

	u := URL("feeds/%s/%s", f.feedType, packageTime)

	httpResp, err := f.client.sendRequest("GET", u, nil, nil)
	if err != nil {
		return nil, err
	}
	defer httpResp.Body.Close()

	switch httpResp.StatusCode {
	case http.StatusBadRequest:
		if resp, err := f.client.parseResponse(httpResp); err != nil {
			if resp.Error.Code == "NotAvailableYet" {
				return nil, errNoAvailableYet
			}
		}
	case http.StatusNotFound:
		return nil, errNotFound
	}

	if httpResp.StatusCode != http.StatusOK {
		return nil, errors.New(httpResp.Status)
	}

	sc := bufio.NewScanner(bzip2.NewReader(httpResp.Body))
	// By default bufio.Scanner uses a buffer that is limited to a maximum size
	// defined by bufio.MaxScanBufferSize (64KB). This is too small for
	// accommodating the large JSONs stored in the feed files. So we create an
	// initial 1MB buffer and let it grow up to 10MB.
	buffer := make([]byte, 1*1024*1024)
	sc.Buffer(buffer, 10*1024*1024)

	objects := make([]*Object, 0)
	for sc.Scan() {
		obj := &Object{}
		if err := json.Unmarshal(sc.Bytes(), obj); err != nil {
			return objects, err
		}
		objects = append(objects, obj)
	}

	return objects, sc.Err()
}

func (f *Feed) retrieve() {
	waitDuration := 20 * time.Second
	missingPackages := 0
loop:
	for {
		packageTime := f.t.Format("200601021504") // YYYYMMDDhhmm
		objects, err := f.getObjects(packageTime)
		objects = objects[f.n:]
		switch err {
		case nil:
			for _, object := range objects {
				if f.sendToChannel(object) == stop {
					break loop
				}
				f.n++
			}
			f.t = f.t.Add(60 * time.Second)
			f.n = 0
			waitDuration = 20 * time.Second
			missingPackages = 0
		case errNoAvailableYet:
			// Feed package is not available yet, let's wait for 1 minute and
			// try again. If Close() is called during the waiting period it
			// exits early and breaks the loop.
			if f.wait(waitDuration) == stop {
				break loop
			}
			waitDuration *= 2
		case errNotFound:
			// The feed tolerates some missing packages, if the number of missing
			// packages is greater than missingPackagesTolerancy an error is
			// returned, if not, it tries to get the next package.
			missingPackages++
			if missingPackages > f.missingPackagesTolerancy {
				f.err = err
				break loop
			}
			f.t = f.t.Add(60 * time.Second)
		default:
			f.err = err
			break loop
		}
	}
	f.stopped = true
	close(f.C)
	close(f.stop)
}
<|endoftext|>


<commit_before>// Package function implements function-level operations.
package function

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/apex/apex/runtime"
	"github.com/apex/apex/shim"
	"github.com/apex/apex/utils"
	"github.com/apex/log"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/aws/aws-sdk-go/service/lambda/lambdaiface"
	"github.com/dustin/go-humanize"
	"github.com/jpillora/archive"
)

// InvocationType determines how an invocation request is made.
type InvocationType string

// Invocation types.
const (
	RequestResponse InvocationType = "RequestResponse"
	Event                          = "Event"
	DryRun                         = "DryRun"
)

// Current alias name.
const CurrentAlias = "current"

// InvokeError records an error from an invocation.
type InvokeError struct {
	Message string   `json:"errorMessage"`
	Type    string   `json:"errorType"`
	Stack   []string `json:"stackTrace"`
	Handled bool
}

// Error message.
func (e *InvokeError) Error() string {
	return e.Message
}

// Config for a Lambda function.
type Config struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	Runtime     string `json:"runtime"`
	Memory      int64  `json:"memory"`
	Timeout     int64  `json:"timeout"`
	Role        string `json:"role"`
}

// Function represents a Lambda function, with configuration loaded
// from the "lambda.json" file on disk. Operations are performed
// against the function directory as the CWD, so os.Chdir() first.
type Function struct {
	Config
	Path    string
	Verbose bool
	Service lambdaiface.LambdaAPI
	Log     log.Interface
	runtime runtime.Runtime
	env     map[string]string
}

// Open the function.json file and prime the config.
func (f *Function) Open() error {
	p, err := os.Open(filepath.Join(f.Path, "function.json"))
	if err != nil {
		return err
	}

	if err := json.NewDecoder(p).Decode(&f.Config); err != nil {
		return err
	}

	r, err := runtime.ByName(f.Runtime)
	if err != nil {
		return err
	}
	f.runtime = r

	return nil
}

// SetEnv sets environment variable `name` to `value`.
func (f *Function) SetEnv(name, value string) {
	if f.env == nil {
		f.env = make(map[string]string)
	}
	f.env[name] = value
}

// Deploy code and then configuration.
func (f *Function) Deploy() error {
	if err := f.DeployCode(); err != nil {
		return err
	}

	return f.DeployConfig()
}

// DeployCode generates a zip and creates or updates the function.
func (f *Function) DeployCode() error {
	f.Log.Info("deploying")

	zip, err := f.ZipBytes()
	if err != nil {
		return err
	}

	info, err := f.Info()

	if e, ok := err.(awserr.Error); ok {
		if e.Code() == "ResourceNotFoundException" {
			return f.Create(zip)
		}
	}

	if err != nil {
		return err
	}

	remoteHash := *info.Configuration.CodeSha256
	localHash := utils.Sha256(zip)

	if localHash == remoteHash {
		f.Log.Info("unchanged")
		return nil
	}

	return f.Update(zip)
}

// DeployConfig deploys changes to configuration.
func (f *Function) DeployConfig() error {
	f.Log.Info("deploying config")

	_, err := f.Service.UpdateFunctionConfiguration(&lambda.UpdateFunctionConfigurationInput{
		FunctionName: &f.Name,
		MemorySize:   &f.Memory,
		Timeout:      &f.Timeout,
		Description:  &f.Description,
		Role:         aws.String(f.Role),
		Handler:      aws.String(f.runtime.Handler()),
	})

	return err
}

// Delete the function including all its versions
func (f *Function) Delete() error {
	f.Log.Info("deleting")
	_, err := f.Service.DeleteFunction(&lambda.DeleteFunctionInput{
		FunctionName: &f.Name,
	})
	return err
}

// Info returns the function information.
func (f *Function) Info() (*lambda.GetFunctionOutput, error) {
	f.Log.Info("fetching config")
	return f.Service.GetFunction(&lambda.GetFunctionInput{
		FunctionName: &f.Name,
	})
}

// Update the function with the given `zip`.
func (f *Function) Update(zip []byte) error {
	f.Log.Info("updating function")

	updated, err := f.Service.UpdateFunctionCode(&lambda.UpdateFunctionCodeInput{
		FunctionName: &f.Name,
		Publish:      aws.Bool(true),
		ZipFile:      zip,
	})

	if err != nil {
		return err
	}

	f.Log.Info("updating alias")

	_, err = f.Service.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    &f.Name,
		Name:            aws.String(CurrentAlias),
		FunctionVersion: updated.Version,
	})

	return err
}

// Create the function with the given `zip`.
func (f *Function) Create(zip []byte) error {
	f.Log.Info("creating function")

	created, err := f.Service.CreateFunction(&lambda.CreateFunctionInput{
		FunctionName: &f.Name,
		Description:  &f.Description,
		MemorySize:   &f.Memory,
		Timeout:      &f.Timeout,
		Runtime:      aws.String(f.runtime.Name()),
		Handler:      aws.String(f.runtime.Handler()),
		Role:         aws.String(f.Role),
		Publish:      aws.Bool(true),
		Code: &lambda.FunctionCode{
			ZipFile: zip,
		},
	})

	if err != nil {
		return err
	}

	f.Log.Info("creating alias")

	_, err = f.Service.CreateAlias(&lambda.CreateAliasInput{
		FunctionName:    &f.Name,
		FunctionVersion: created.Version,
		Name:            aws.String(CurrentAlias),
	})

	return err
}

// Invoke the remote Lambda function, returning the response and logs, if any.
func (f *Function) Invoke(event, context interface{}, kind InvocationType) (reply, logs io.Reader, err error) {
	eventBytes, err := json.Marshal(event)
	if err != nil {
		return nil, nil, err
	}

	contextBytes, err := json.Marshal(context)
	if err != nil {
		return nil, nil, err
	}

	res, err := f.Service.Invoke(&lambda.InvokeInput{
		ClientContext:  aws.String(base64.StdEncoding.EncodeToString(contextBytes)),
		FunctionName:   aws.String(f.Name),
		InvocationType: aws.String(string(kind)),
		LogType:        aws.String("Tail"),
		Qualifier:      aws.String("current"),
		Payload:        eventBytes,
	})

	if err != nil {
		return nil, nil, err
	}

	if res.FunctionError != nil {
		e := &InvokeError{
			Handled: *res.FunctionError == "Handled",
		}

		if err := json.Unmarshal(res.Payload, e); err != nil {
			return nil, nil, err
		}

		return nil, nil, e
	}

	if kind == Event {
		return bytes.NewReader(nil), bytes.NewReader(nil), nil
	}

	logs = base64.NewDecoder(base64.StdEncoding, strings.NewReader(*res.LogResult))
	reply = bytes.NewReader(res.Payload)
	return reply, logs, nil
}

// Rollback the function to the previous or specified version.
func (f *Function) Rollback(version ...string) error {
	f.Log.Info("rolling back")

	isVersionSpecified := len(version) > 0
	var specifiedVersion string
	if isVersionSpecified {
		specifiedVersion = version[0]
	}

	alias, err := f.Service.GetAlias(&lambda.GetAliasInput{
		FunctionName: &f.Name,
		Name:         aws.String(CurrentAlias),
	})

	if err != nil {
		return err
	}

	f.Log.Infof("current version: %s", *alias.FunctionVersion)

	if isVersionSpecified && specifiedVersion == *alias.FunctionVersion {
		return errors.New("Specified version currently deployed.")
	}

	list, err := f.Service.ListVersionsByFunction(&lambda.ListVersionsByFunctionInput{
		FunctionName: &f.Name,
	})

	if err != nil {
		return err
	}

	versions := list.Versions
	_, versions = versions[0], versions[1:] // remove $LATEST
	if len(versions) < 2 {
		return errors.New("Can't rollback. Only one version deployed.")
	}
	latestVersion := *versions[len(versions)-1].Version
	prevVersion := *versions[len(versions)-2].Version

	rollbackToVersion := latestVersion
	if isVersionSpecified {
		rollbackToVersion = specifiedVersion
	} else if *alias.FunctionVersion == latestVersion {
		rollbackToVersion = prevVersion
	}

	f.Log.Infof("rollback to version: %s", rollbackToVersion)

	_, err = f.Service.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    &f.Name,
		Name:            aws.String(CurrentAlias),
		FunctionVersion: &rollbackToVersion,
	})

	return err
}

// Clean removes build artifacts from compiled runtimes.
func (f *Function) Clean() error {
	if r, ok := f.runtime.(runtime.CompiledRuntime); ok {
		return r.Clean(f.Path)
	}
	return nil
}

// Zip returns the zipped contents of the function.
func (f *Function) Zip() (io.Reader, error) {
	buf := new(bytes.Buffer)
	zip := archive.NewZipWriter(buf)

	if r, ok := f.runtime.(runtime.CompiledRuntime); ok {
		f.Log.Debugf("compiling")
		if err := r.Build(f.Path); err != nil {
			return nil, fmt.Errorf("compiling: %s", err)
		}
	}

	// TODO(tj): remove or add --env flag back
	if f.env != nil {
		f.Log.Debugf("adding .env.json")

		b, err := json.Marshal(f.env)
		if err != nil {
			return nil, err
		}

		zip.AddBytes(".env.json", b)
	}

	if f.runtime.Shimmed() {
		f.Log.Debugf("adding nodejs shim")
		zip.AddBytes("index.js", shim.MustAsset("index.js"))
		zip.AddBytes("byline.js", shim.MustAsset("byline.js"))
	}

	if err := zip.AddDir(f.Path); err != nil {
		return nil, err
	}

	if err := zip.Close(); err != nil {
		return nil, err
	}

	return buf, nil
}

// ZipBytes returns the generated zip as bytes.
func (f *Function) ZipBytes() ([]byte, error) {
	f.Log.Debugf("creating zip")

	r, err := f.Zip()
	if err != nil {
		return nil, err
	}

	b, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}

	f.Log.Infof("created zip (%s)", humanize.Bytes(uint64(len(b))))
	return b, nil
}
<commit_msg>refactor to use CurrentAlias<commit_after>// Package function implements function-level operations.
package function

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/apex/apex/runtime"
	"github.com/apex/apex/shim"
	"github.com/apex/apex/utils"
	"github.com/apex/log"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/aws/aws-sdk-go/service/lambda/lambdaiface"
	"github.com/dustin/go-humanize"
	"github.com/jpillora/archive"
)

// InvocationType determines how an invocation request is made.
type InvocationType string

// Invocation types.
const (
	RequestResponse InvocationType = "RequestResponse"
	Event                          = "Event"
	DryRun                         = "DryRun"
)

// Current alias name.
const CurrentAlias = "current"

// InvokeError records an error from an invocation.
type InvokeError struct {
	Message string   `json:"errorMessage"`
	Type    string   `json:"errorType"`
	Stack   []string `json:"stackTrace"`
	Handled bool
}

// Error message.
func (e *InvokeError) Error() string {
	return e.Message
}

// Config for a Lambda function.
type Config struct {
	Name        string `json:"name"`
	Description string `json:"description"`
	Runtime     string `json:"runtime"`
	Memory      int64  `json:"memory"`
	Timeout     int64  `json:"timeout"`
	Role        string `json:"role"`
}

// Function represents a Lambda function, with configuration loaded
// from the "lambda.json" file on disk. Operations are performed
// against the function directory as the CWD, so os.Chdir() first.
type Function struct {
	Config
	Path    string
	Verbose bool
	Service lambdaiface.LambdaAPI
	Log     log.Interface
	runtime runtime.Runtime
	env     map[string]string
}

// Open the function.json file and prime the config.
func (f *Function) Open() error {
	p, err := os.Open(filepath.Join(f.Path, "function.json"))
	if err != nil {
		return err
	}

	if err := json.NewDecoder(p).Decode(&f.Config); err != nil {
		return err
	}

	r, err := runtime.ByName(f.Runtime)
	if err != nil {
		return err
	}
	f.runtime = r

	return nil
}

// SetEnv sets environment variable `name` to `value`.
func (f *Function) SetEnv(name, value string) {
	if f.env == nil {
		f.env = make(map[string]string)
	}
	f.env[name] = value
}

// Deploy code and then configuration.
func (f *Function) Deploy() error {
	if err := f.DeployCode(); err != nil {
		return err
	}

	return f.DeployConfig()
}

// DeployCode generates a zip and creates or updates the function.
func (f *Function) DeployCode() error {
	f.Log.Info("deploying")

	zip, err := f.ZipBytes()
	if err != nil {
		return err
	}

	info, err := f.Info()

	if e, ok := err.(awserr.Error); ok {
		if e.Code() == "ResourceNotFoundException" {
			return f.Create(zip)
		}
	}

	if err != nil {
		return err
	}

	remoteHash := *info.Configuration.CodeSha256
	localHash := utils.Sha256(zip)

	if localHash == remoteHash {
		f.Log.Info("unchanged")
		return nil
	}

	return f.Update(zip)
}

// DeployConfig deploys changes to configuration.
func (f *Function) DeployConfig() error {
	f.Log.Info("deploying config")

	_, err := f.Service.UpdateFunctionConfiguration(&lambda.UpdateFunctionConfigurationInput{
		FunctionName: &f.Name,
		MemorySize:   &f.Memory,
		Timeout:      &f.Timeout,
		Description:  &f.Description,
		Role:         aws.String(f.Role),
		Handler:      aws.String(f.runtime.Handler()),
	})

	return err
}

// Delete the function including all its versions
func (f *Function) Delete() error {
	f.Log.Info("deleting")
	_, err := f.Service.DeleteFunction(&lambda.DeleteFunctionInput{
		FunctionName: &f.Name,
	})
	return err
}

// Info returns the function information.
func (f *Function) Info() (*lambda.GetFunctionOutput, error) {
	f.Log.Info("fetching config")
	return f.Service.GetFunction(&lambda.GetFunctionInput{
		FunctionName: &f.Name,
	})
}

// Update the function with the given `zip`.
func (f *Function) Update(zip []byte) error {
	f.Log.Info("updating function")

	updated, err := f.Service.UpdateFunctionCode(&lambda.UpdateFunctionCodeInput{
		FunctionName: &f.Name,
		Publish:      aws.Bool(true),
		ZipFile:      zip,
	})

	if err != nil {
		return err
	}

	f.Log.Info("updating alias")

	_, err = f.Service.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    &f.Name,
		Name:            aws.String(CurrentAlias),
		FunctionVersion: updated.Version,
	})

	return err
}

// Create the function with the given `zip`.
func (f *Function) Create(zip []byte) error {
	f.Log.Info("creating function")

	created, err := f.Service.CreateFunction(&lambda.CreateFunctionInput{
		FunctionName: &f.Name,
		Description:  &f.Description,
		MemorySize:   &f.Memory,
		Timeout:      &f.Timeout,
		Runtime:      aws.String(f.runtime.Name()),
		Handler:      aws.String(f.runtime.Handler()),
		Role:         aws.String(f.Role),
		Publish:      aws.Bool(true),
		Code: &lambda.FunctionCode{
			ZipFile: zip,
		},
	})

	if err != nil {
		return err
	}

	f.Log.Info("creating alias")

	_, err = f.Service.CreateAlias(&lambda.CreateAliasInput{
		FunctionName:    &f.Name,
		FunctionVersion: created.Version,
		Name:            aws.String(CurrentAlias),
	})

	return err
}

// Invoke the remote Lambda function, returning the response and logs, if any.
func (f *Function) Invoke(event, context interface{}, kind InvocationType) (reply, logs io.Reader, err error) {
	eventBytes, err := json.Marshal(event)
	if err != nil {
		return nil, nil, err
	}

	contextBytes, err := json.Marshal(context)
	if err != nil {
		return nil, nil, err
	}

	res, err := f.Service.Invoke(&lambda.InvokeInput{
		ClientContext:  aws.String(base64.StdEncoding.EncodeToString(contextBytes)),
		FunctionName:   aws.String(f.Name),
		InvocationType: aws.String(string(kind)),
		LogType:        aws.String("Tail"),
		Qualifier:      aws.String(CurrentAlias),
		Payload:        eventBytes,
	})

	if err != nil {
		return nil, nil, err
	}

	if res.FunctionError != nil {
		e := &InvokeError{
			Handled: *res.FunctionError == "Handled",
		}

		if err := json.Unmarshal(res.Payload, e); err != nil {
			return nil, nil, err
		}

		return nil, nil, e
	}

	if kind == Event {
		return bytes.NewReader(nil), bytes.NewReader(nil), nil
	}

	logs = base64.NewDecoder(base64.StdEncoding, strings.NewReader(*res.LogResult))
	reply = bytes.NewReader(res.Payload)
	return reply, logs, nil
}

// Rollback the function to the previous or specified version.
func (f *Function) Rollback(version ...string) error {
	f.Log.Info("rolling back")

	isVersionSpecified := len(version) > 0
	var specifiedVersion string
	if isVersionSpecified {
		specifiedVersion = version[0]
	}

	alias, err := f.Service.GetAlias(&lambda.GetAliasInput{
		FunctionName: &f.Name,
		Name:         aws.String(CurrentAlias),
	})

	if err != nil {
		return err
	}

	f.Log.Infof("current version: %s", *alias.FunctionVersion)

	if isVersionSpecified && specifiedVersion == *alias.FunctionVersion {
		return errors.New("Specified version currently deployed.")
	}

	list, err := f.Service.ListVersionsByFunction(&lambda.ListVersionsByFunctionInput{
		FunctionName: &f.Name,
	})

	if err != nil {
		return err
	}

	versions := list.Versions
	_, versions = versions[0], versions[1:] // remove $LATEST
	if len(versions) < 2 {
		return errors.New("Can't rollback. Only one version deployed.")
	}
	latestVersion := *versions[len(versions)-1].Version
	prevVersion := *versions[len(versions)-2].Version

	rollbackToVersion := latestVersion
	if isVersionSpecified {
		rollbackToVersion = specifiedVersion
	} else if *alias.FunctionVersion == latestVersion {
		rollbackToVersion = prevVersion
	}

	f.Log.Infof("rollback to version: %s", rollbackToVersion)

	_, err = f.Service.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    &f.Name,
		Name:            aws.String(CurrentAlias),
		FunctionVersion: &rollbackToVersion,
	})

	return err
}

// Clean removes build artifacts from compiled runtimes.
func (f *Function) Clean() error {
	if r, ok := f.runtime.(runtime.CompiledRuntime); ok {
		return r.Clean(f.Path)
	}
	return nil
}

// Zip returns the zipped contents of the function.
func (f *Function) Zip() (io.Reader, error) {
	buf := new(bytes.Buffer)
	zip := archive.NewZipWriter(buf)

	if r, ok := f.runtime.(runtime.CompiledRuntime); ok {
		f.Log.Debugf("compiling")
		if err := r.Build(f.Path); err != nil {
			return nil, fmt.Errorf("compiling: %s", err)
		}
	}

	// TODO(tj): remove or add --env flag back
	if f.env != nil {
		f.Log.Debugf("adding .env.json")

		b, err := json.Marshal(f.env)
		if err != nil {
			return nil, err
		}

		zip.AddBytes(".env.json", b)
	}

	if f.runtime.Shimmed() {
		f.Log.Debugf("adding nodejs shim")
		zip.AddBytes("index.js", shim.MustAsset("index.js"))
		zip.AddBytes("byline.js", shim.MustAsset("byline.js"))
	}

	if err := zip.AddDir(f.Path); err != nil {
		return nil, err
	}

	if err := zip.Close(); err != nil {
		return nil, err
	}

	return buf, nil
}

// ZipBytes returns the generated zip as bytes.
func (f *Function) ZipBytes() ([]byte, error) {
	f.Log.Debugf("creating zip")

	r, err := f.Zip()
	if err != nil {
		return nil, err
	}

	b, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}

	f.Log.Infof("created zip (%s)", humanize.Bytes(uint64(len(b))))
	return b, nil
}
<|endoftext|>


<commit_before>// Package function implements function-level operations.
package function

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/aws/aws-sdk-go/service/lambda/lambdaiface"
	"github.com/dustin/go-humanize"
	"github.com/jpillora/archive"
	"gopkg.in/validator.v2"

	"github.com/apex/apex/hooks"
	"github.com/apex/apex/runtime"
	"github.com/apex/apex/shim"
	"github.com/apex/apex/utils"
	"github.com/apex/log"
)

// InvocationType determines how an invocation request is made.
type InvocationType string

// Invocation types.
const (
	RequestResponse InvocationType = "RequestResponse"
	Event                          = "Event"
	DryRun                         = "DryRun"
)

// Current alias name.
const CurrentAlias = "current"

// InvokeError records an error from an invocation.
type InvokeError struct {
	Message string   `json:"errorMessage"`
	Type    string   `json:"errorType"`
	Stack   []string `json:"stackTrace"`
	Handled bool
}

// Error message.
func (e *InvokeError) Error() string {
	return e.Message
}

// Config for a Lambda function.
type Config struct {
	Description string      `json:"description"`
	Runtime     string      `json:"runtime" validate:"nonzero"`
	Memory      int64       `json:"memory" validate:"nonzero"`
	Timeout     int64       `json:"timeout" validate:"nonzero"`
	Role        string      `json:"role" validate:"nonzero"`
	Hooks       hooks.Hooks `json:"hooks"`
}

// Function represents a Lambda function, with configuration loaded
// from the "function.json" file on disk. Operations are performed
// against the function directory as the CWD, so os.Chdir() first.
type Function struct {
	Config
	Name            string
	FunctionName    string
	Path            string
	Service         lambdaiface.LambdaAPI
	Log             log.Interface
	IgnoredPatterns []string
	runtime         runtime.Runtime
	env             map[string]string
}

// Open the function.json file and prime the config.
func (f *Function) Open() error {
	p, err := os.Open(filepath.Join(f.Path, "function.json"))
	if err == nil {
		if err := json.NewDecoder(p).Decode(&f.Config); err != nil {
			return err
		}
	}

	if f.Runtime == "" {
		if runtimeName, err := runtime.Detect(f.Path); err == nil {
			f.Runtime = runtimeName
		}
	}

	if err := validator.Validate(&f.Config); err != nil {
		return fmt.Errorf("error opening function %s: %s", f.Name, err.Error())
	}

	r, err := runtime.ByName(f.Runtime)
	if err != nil {
		return err
	}
	f.runtime = r

	patterns, err := utils.ReadIgnoreFile(f.Path)
	if err != nil {
		return err
	}
	f.IgnoredPatterns = append(f.IgnoredPatterns, patterns...)

	f.Log = f.Log.WithField("function", f.Name)

	return nil
}

// SetEnv sets environment variable `name` to `value`.
func (f *Function) SetEnv(name, value string) {
	if f.env == nil {
		f.env = make(map[string]string)
	}
	f.env[name] = value
}

// Deploy code and then configuration.
func (f *Function) Deploy() error {
	if err := f.DeployCode(); err != nil {
		return err
	}

	return f.DeployConfig()
}

// DeployCode generates a zip and creates or updates the function.
func (f *Function) DeployCode() error {
	f.Log.Info("deploying")

	zip, err := f.BuildBytes()
	if err != nil {
		return err
	}

	config, err := f.GetConfig()

	if e, ok := err.(awserr.Error); ok {
		if e.Code() == "ResourceNotFoundException" {
			return f.Create(zip)
		}
	}

	if err != nil {
		return err
	}

	remoteHash := *config.Configuration.CodeSha256
	localHash := utils.Sha256(zip)

	if localHash == remoteHash {
		f.Log.Info("unchanged")
		return nil
	}

	return f.Update(zip)
}

// DeployConfig deploys changes to configuration.
func (f *Function) DeployConfig() error {
	f.Log.Info("deploying config")

	_, err := f.Service.UpdateFunctionConfiguration(&lambda.UpdateFunctionConfigurationInput{
		FunctionName: &f.FunctionName,
		MemorySize:   &f.Memory,
		Timeout:      &f.Timeout,
		Description:  &f.Description,
		Role:         aws.String(f.Role),
		Handler:      aws.String(f.runtime.Handler()),
	})

	return err
}

// Delete the function including all its versions
func (f *Function) Delete() error {
	f.Log.Info("deleting")
	_, err := f.Service.DeleteFunction(&lambda.DeleteFunctionInput{
		FunctionName: &f.FunctionName,
	})
	return err
}

// GetConfig returns the function configuration.
func (f *Function) GetConfig() (*lambda.GetFunctionOutput, error) {
	f.Log.Debug("fetching config")
	return f.Service.GetFunction(&lambda.GetFunctionInput{
		FunctionName: &f.FunctionName,
	})
}

// Update the function with the given `zip`.
func (f *Function) Update(zip []byte) error {
	f.Log.Info("updating function")

	updated, err := f.Service.UpdateFunctionCode(&lambda.UpdateFunctionCodeInput{
		FunctionName: &f.FunctionName,
		Publish:      aws.Bool(true),
		ZipFile:      zip,
	})

	if err != nil {
		return err
	}

	f.Log.Info("updating alias")

	_, err = f.Service.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    &f.FunctionName,
		Name:            aws.String(CurrentAlias),
		FunctionVersion: updated.Version,
	})

	if err != nil {
		return nil
	}

	f.Log.WithFields(log.Fields{
		"version": *updated.Version,
		"name":    f.FunctionName,
	}).Info("deployed")

	return nil
}

// Create the function with the given `zip`.
func (f *Function) Create(zip []byte) error {
	f.Log.Info("creating function")

	created, err := f.Service.CreateFunction(&lambda.CreateFunctionInput{
		FunctionName: &f.FunctionName,
		Description:  &f.Description,
		MemorySize:   &f.Memory,
		Timeout:      &f.Timeout,
		Runtime:      aws.String(f.runtime.Name()),
		Handler:      aws.String(f.runtime.Handler()),
		Role:         aws.String(f.Role),
		Publish:      aws.Bool(true),
		Code: &lambda.FunctionCode{
			ZipFile: zip,
		},
	})

	if err != nil {
		return err
	}

	f.Log.Info("creating alias")

	_, err = f.Service.CreateAlias(&lambda.CreateAliasInput{
		FunctionName:    &f.FunctionName,
		FunctionVersion: created.Version,
		Name:            aws.String(CurrentAlias),
	})

	if err != nil {
		return nil
	}

	f.Log.WithFields(log.Fields{
		"version": *created.Version,
		"name":    f.FunctionName,
	}).Info("deployed")

	return nil
}

// Invoke the remote Lambda function, returning the response and logs, if any.
func (f *Function) Invoke(event, context interface{}, kind InvocationType) (reply, logs io.Reader, err error) {
	eventBytes, err := json.Marshal(event)
	if err != nil {
		return nil, nil, err
	}

	contextBytes, err := json.Marshal(context)
	if err != nil {
		return nil, nil, err
	}

	res, err := f.Service.Invoke(&lambda.InvokeInput{
		ClientContext:  aws.String(base64.StdEncoding.EncodeToString(contextBytes)),
		FunctionName:   &f.FunctionName,
		InvocationType: aws.String(string(kind)),
		LogType:        aws.String("Tail"),
		Qualifier:      aws.String(CurrentAlias),
		Payload:        eventBytes,
	})

	if err != nil {
		return nil, nil, err
	}

	if res.FunctionError != nil {
		e := &InvokeError{
			Handled: *res.FunctionError == "Handled",
		}

		if err := json.Unmarshal(res.Payload, e); err != nil {
			return nil, nil, err
		}

		return nil, nil, e
	}

	if kind == Event {
		return bytes.NewReader(nil), bytes.NewReader(nil), nil
	}

	logs = base64.NewDecoder(base64.StdEncoding, strings.NewReader(*res.LogResult))
	reply = bytes.NewReader(res.Payload)
	return reply, logs, nil
}

// Rollback the function to the previous.
func (f *Function) Rollback() error {
	f.Log.Info("rolling back")

	alias, err := f.Service.GetAlias(&lambda.GetAliasInput{
		FunctionName: &f.FunctionName,
		Name:         aws.String(CurrentAlias),
	})

	if err != nil {
		return err
	}

	f.Log.Infof("current version: %s", *alias.FunctionVersion)

	list, err := f.Service.ListVersionsByFunction(&lambda.ListVersionsByFunctionInput{
		FunctionName: &f.FunctionName,
	})

	if err != nil {
		return err
	}

	versions := list.Versions[1:] // remove $LATEST
	if len(versions) < 2 {
		return errors.New("Can't rollback. Only one version deployed.")
	}

	latest := *versions[len(versions)-1].Version
	prev := *versions[len(versions)-2].Version
	rollback := latest

	if *alias.FunctionVersion == latest {
		rollback = prev
	}

	f.Log.Infof("rollback to version: %s", rollback)

	_, err = f.Service.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    &f.FunctionName,
		Name:            aws.String(CurrentAlias),
		FunctionVersion: &rollback,
	})

	return err
}

// RollbackVersion the function to the specified version.
func (f *Function) RollbackVersion(version string) error {
	f.Log.Info("rolling back")

	alias, err := f.Service.GetAlias(&lambda.GetAliasInput{
		FunctionName: &f.FunctionName,
		Name:         aws.String(CurrentAlias),
	})

	if err != nil {
		return err
	}

	f.Log.Infof("current version: %s", *alias.FunctionVersion)

	if version == *alias.FunctionVersion {
		return errors.New("Specified version currently deployed.")
	}

	_, err = f.Service.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    &f.FunctionName,
		Name:            aws.String(CurrentAlias),
		FunctionVersion: &version,
	})

	return err
}

// BuildBytes returns the generated zip as bytes.
func (f *Function) BuildBytes() ([]byte, error) {
	r, err := f.Build()
	if err != nil {
		return nil, err
	}

	b, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}

	f.Log.Infof("created build (%s)", humanize.Bytes(uint64(len(b))))
	return b, nil
}

// Build returns the zipped contents of the function.
func (f *Function) Build() (io.Reader, error) {
	f.Log.Debugf("creating build")

	if err := f.RunHook("build"); err != nil {
		return nil, err
	}

	buf := new(bytes.Buffer)
	zip := archive.NewZipWriter(buf)

	if r, ok := f.runtime.(runtime.CompiledRuntime); ok {
		f.Log.Debugf("compiling")
		if err := r.Build(f.Path); err != nil {
			return nil, fmt.Errorf("compiling: %s", err)
		}
	}

	if f.env != nil {
		f.Log.Debugf("adding .env.json")

		b, err := json.Marshal(f.env)
		if err != nil {
			return nil, err
		}

		zip.AddBytes(".env.json", b)
	}

	if f.runtime.Shimmed() {
		f.Log.Debugf("adding nodejs shim")
		zip.AddBytes("index.js", shim.MustAsset("index.js"))
		zip.AddBytes("byline.js", shim.MustAsset("byline.js"))
	}

	files, err := utils.LoadFiles(f.Path, f.IgnoredPatterns)
	if err != nil {
		return nil, err
	}

	for path, file := range files {
		f.Log.WithField("file", path).Debug("add file")
		if err := zip.AddFile(path, file); err != nil {
			return nil, err
		}
		defer file.Close()
	}

	if err := zip.Close(); err != nil {
		return nil, err
	}

	return buf, nil
}

// Clean removes build artifacts from compiled runtimes.
func (f *Function) Clean() error {
	if err := f.RunHook("clean"); err != nil {
		return err
	}

	if r, ok := f.runtime.(runtime.CompiledRuntime); ok {
		return r.Clean(f.Path)
	}

	return nil
}

// A HookError represents a failed hook command.
type HookError struct {
	Name    string
	Command string
	Output  string
}

// Error string.
func (e *HookError) Error() string {
	return fmt.Sprintf("hook %q: %s", e.Name, e.Output)
}

// RunHook executes hook `name` in the function's directory.
func (f *Function) RunHook(name string) error {
	var command string

	switch name {
	case "clean":
		command = f.Hooks.Clean
	case "build":
		command = f.Hooks.Build
	}

	if command == "" {
		return nil
	}

	f.Log.WithFields(log.Fields{
		"hook":    name,
		"command": command,
	}).Debug("hook")

	cmd := exec.Command("sh", "-c", command)
	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, fmt.Sprintf("FUNCTION=%s", f.Name))
	cmd.Dir = f.Path

	b, err := cmd.CombinedOutput()
	if err != nil {
		return &HookError{
			Name:    name,
			Command: command,
			Output:  string(b),
		}
	}

	return nil
}
<commit_msg>remove deferring of file Close() for builds<commit_after>// Package function implements function-level operations.
package function

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/aws/aws-sdk-go/service/lambda/lambdaiface"
	"github.com/dustin/go-humanize"
	"github.com/jpillora/archive"
	"gopkg.in/validator.v2"

	"github.com/apex/apex/hooks"
	"github.com/apex/apex/runtime"
	"github.com/apex/apex/shim"
	"github.com/apex/apex/utils"
	"github.com/apex/log"
)

// InvocationType determines how an invocation request is made.
type InvocationType string

// Invocation types.
const (
	RequestResponse InvocationType = "RequestResponse"
	Event                          = "Event"
	DryRun                         = "DryRun"
)

// Current alias name.
const CurrentAlias = "current"

// InvokeError records an error from an invocation.
type InvokeError struct {
	Message string   `json:"errorMessage"`
	Type    string   `json:"errorType"`
	Stack   []string `json:"stackTrace"`
	Handled bool
}

// Error message.
func (e *InvokeError) Error() string {
	return e.Message
}

// Config for a Lambda function.
type Config struct {
	Description string      `json:"description"`
	Runtime     string      `json:"runtime" validate:"nonzero"`
	Memory      int64       `json:"memory" validate:"nonzero"`
	Timeout     int64       `json:"timeout" validate:"nonzero"`
	Role        string      `json:"role" validate:"nonzero"`
	Hooks       hooks.Hooks `json:"hooks"`
}

// Function represents a Lambda function, with configuration loaded
// from the "function.json" file on disk. Operations are performed
// against the function directory as the CWD, so os.Chdir() first.
type Function struct {
	Config
	Name            string
	FunctionName    string
	Path            string
	Service         lambdaiface.LambdaAPI
	Log             log.Interface
	IgnoredPatterns []string
	runtime         runtime.Runtime
	env             map[string]string
}

// Open the function.json file and prime the config.
func (f *Function) Open() error {
	p, err := os.Open(filepath.Join(f.Path, "function.json"))
	if err == nil {
		if err := json.NewDecoder(p).Decode(&f.Config); err != nil {
			return err
		}
	}

	if f.Runtime == "" {
		if runtimeName, err := runtime.Detect(f.Path); err == nil {
			f.Runtime = runtimeName
		}
	}

	if err := validator.Validate(&f.Config); err != nil {
		return fmt.Errorf("error opening function %s: %s", f.Name, err.Error())
	}

	r, err := runtime.ByName(f.Runtime)
	if err != nil {
		return err
	}
	f.runtime = r

	patterns, err := utils.ReadIgnoreFile(f.Path)
	if err != nil {
		return err
	}
	f.IgnoredPatterns = append(f.IgnoredPatterns, patterns...)

	f.Log = f.Log.WithField("function", f.Name)

	return nil
}

// SetEnv sets environment variable `name` to `value`.
func (f *Function) SetEnv(name, value string) {
	if f.env == nil {
		f.env = make(map[string]string)
	}
	f.env[name] = value
}

// Deploy code and then configuration.
func (f *Function) Deploy() error {
	if err := f.DeployCode(); err != nil {
		return err
	}

	return f.DeployConfig()
}

// DeployCode generates a zip and creates or updates the function.
func (f *Function) DeployCode() error {
	f.Log.Info("deploying")

	zip, err := f.BuildBytes()
	if err != nil {
		return err
	}

	config, err := f.GetConfig()

	if e, ok := err.(awserr.Error); ok {
		if e.Code() == "ResourceNotFoundException" {
			return f.Create(zip)
		}
	}

	if err != nil {
		return err
	}

	remoteHash := *config.Configuration.CodeSha256
	localHash := utils.Sha256(zip)

	if localHash == remoteHash {
		f.Log.Info("unchanged")
		return nil
	}

	return f.Update(zip)
}

// DeployConfig deploys changes to configuration.
func (f *Function) DeployConfig() error {
	f.Log.Info("deploying config")

	_, err := f.Service.UpdateFunctionConfiguration(&lambda.UpdateFunctionConfigurationInput{
		FunctionName: &f.FunctionName,
		MemorySize:   &f.Memory,
		Timeout:      &f.Timeout,
		Description:  &f.Description,
		Role:         aws.String(f.Role),
		Handler:      aws.String(f.runtime.Handler()),
	})

	return err
}

// Delete the function including all its versions
func (f *Function) Delete() error {
	f.Log.Info("deleting")
	_, err := f.Service.DeleteFunction(&lambda.DeleteFunctionInput{
		FunctionName: &f.FunctionName,
	})
	return err
}

// GetConfig returns the function configuration.
func (f *Function) GetConfig() (*lambda.GetFunctionOutput, error) {
	f.Log.Debug("fetching config")
	return f.Service.GetFunction(&lambda.GetFunctionInput{
		FunctionName: &f.FunctionName,
	})
}

// Update the function with the given `zip`.
func (f *Function) Update(zip []byte) error {
	f.Log.Info("updating function")

	updated, err := f.Service.UpdateFunctionCode(&lambda.UpdateFunctionCodeInput{
		FunctionName: &f.FunctionName,
		Publish:      aws.Bool(true),
		ZipFile:      zip,
	})

	if err != nil {
		return err
	}

	f.Log.Info("updating alias")

	_, err = f.Service.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    &f.FunctionName,
		Name:            aws.String(CurrentAlias),
		FunctionVersion: updated.Version,
	})

	if err != nil {
		return nil
	}

	f.Log.WithFields(log.Fields{
		"version": *updated.Version,
		"name":    f.FunctionName,
	}).Info("deployed")

	return nil
}

// Create the function with the given `zip`.
func (f *Function) Create(zip []byte) error {
	f.Log.Info("creating function")

	created, err := f.Service.CreateFunction(&lambda.CreateFunctionInput{
		FunctionName: &f.FunctionName,
		Description:  &f.Description,
		MemorySize:   &f.Memory,
		Timeout:      &f.Timeout,
		Runtime:      aws.String(f.runtime.Name()),
		Handler:      aws.String(f.runtime.Handler()),
		Role:         aws.String(f.Role),
		Publish:      aws.Bool(true),
		Code: &lambda.FunctionCode{
			ZipFile: zip,
		},
	})

	if err != nil {
		return err
	}

	f.Log.Info("creating alias")

	_, err = f.Service.CreateAlias(&lambda.CreateAliasInput{
		FunctionName:    &f.FunctionName,
		FunctionVersion: created.Version,
		Name:            aws.String(CurrentAlias),
	})

	if err != nil {
		return nil
	}

	f.Log.WithFields(log.Fields{
		"version": *created.Version,
		"name":    f.FunctionName,
	}).Info("deployed")

	return nil
}

// Invoke the remote Lambda function, returning the response and logs, if any.
func (f *Function) Invoke(event, context interface{}, kind InvocationType) (reply, logs io.Reader, err error) {
	eventBytes, err := json.Marshal(event)
	if err != nil {
		return nil, nil, err
	}

	contextBytes, err := json.Marshal(context)
	if err != nil {
		return nil, nil, err
	}

	res, err := f.Service.Invoke(&lambda.InvokeInput{
		ClientContext:  aws.String(base64.StdEncoding.EncodeToString(contextBytes)),
		FunctionName:   &f.FunctionName,
		InvocationType: aws.String(string(kind)),
		LogType:        aws.String("Tail"),
		Qualifier:      aws.String(CurrentAlias),
		Payload:        eventBytes,
	})

	if err != nil {
		return nil, nil, err
	}

	if res.FunctionError != nil {
		e := &InvokeError{
			Handled: *res.FunctionError == "Handled",
		}

		if err := json.Unmarshal(res.Payload, e); err != nil {
			return nil, nil, err
		}

		return nil, nil, e
	}

	if kind == Event {
		return bytes.NewReader(nil), bytes.NewReader(nil), nil
	}

	logs = base64.NewDecoder(base64.StdEncoding, strings.NewReader(*res.LogResult))
	reply = bytes.NewReader(res.Payload)
	return reply, logs, nil
}

// Rollback the function to the previous.
func (f *Function) Rollback() error {
	f.Log.Info("rolling back")

	alias, err := f.Service.GetAlias(&lambda.GetAliasInput{
		FunctionName: &f.FunctionName,
		Name:         aws.String(CurrentAlias),
	})

	if err != nil {
		return err
	}

	f.Log.Infof("current version: %s", *alias.FunctionVersion)

	list, err := f.Service.ListVersionsByFunction(&lambda.ListVersionsByFunctionInput{
		FunctionName: &f.FunctionName,
	})

	if err != nil {
		return err
	}

	versions := list.Versions[1:] // remove $LATEST
	if len(versions) < 2 {
		return errors.New("Can't rollback. Only one version deployed.")
	}

	latest := *versions[len(versions)-1].Version
	prev := *versions[len(versions)-2].Version
	rollback := latest

	if *alias.FunctionVersion == latest {
		rollback = prev
	}

	f.Log.Infof("rollback to version: %s", rollback)

	_, err = f.Service.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    &f.FunctionName,
		Name:            aws.String(CurrentAlias),
		FunctionVersion: &rollback,
	})

	return err
}

// RollbackVersion the function to the specified version.
func (f *Function) RollbackVersion(version string) error {
	f.Log.Info("rolling back")

	alias, err := f.Service.GetAlias(&lambda.GetAliasInput{
		FunctionName: &f.FunctionName,
		Name:         aws.String(CurrentAlias),
	})

	if err != nil {
		return err
	}

	f.Log.Infof("current version: %s", *alias.FunctionVersion)

	if version == *alias.FunctionVersion {
		return errors.New("Specified version currently deployed.")
	}

	_, err = f.Service.UpdateAlias(&lambda.UpdateAliasInput{
		FunctionName:    &f.FunctionName,
		Name:            aws.String(CurrentAlias),
		FunctionVersion: &version,
	})

	return err
}

// BuildBytes returns the generated zip as bytes.
func (f *Function) BuildBytes() ([]byte, error) {
	r, err := f.Build()
	if err != nil {
		return nil, err
	}

	b, err := ioutil.ReadAll(r)
	if err != nil {
		return nil, err
	}

	f.Log.Infof("created build (%s)", humanize.Bytes(uint64(len(b))))
	return b, nil
}

// Build returns the zipped contents of the function.
func (f *Function) Build() (io.Reader, error) {
	f.Log.Debugf("creating build")

	if err := f.RunHook("build"); err != nil {
		return nil, err
	}

	buf := new(bytes.Buffer)
	zip := archive.NewZipWriter(buf)

	if r, ok := f.runtime.(runtime.CompiledRuntime); ok {
		f.Log.Debugf("compiling")
		if err := r.Build(f.Path); err != nil {
			return nil, fmt.Errorf("compiling: %s", err)
		}
	}

	if f.env != nil {
		f.Log.Debugf("adding .env.json")

		b, err := json.Marshal(f.env)
		if err != nil {
			return nil, err
		}

		zip.AddBytes(".env.json", b)
	}

	if f.runtime.Shimmed() {
		f.Log.Debugf("adding nodejs shim")
		zip.AddBytes("index.js", shim.MustAsset("index.js"))
		zip.AddBytes("byline.js", shim.MustAsset("byline.js"))
	}

	files, err := utils.LoadFiles(f.Path, f.IgnoredPatterns)
	if err != nil {
		return nil, err
	}

	for path, file := range files {
		f.Log.WithField("file", path).Debug("add file")
		if err := zip.AddFile(path, file); err != nil {
			return nil, err
		}

		if err := file.Close(); err != nil {
			return nil, err
		}
	}

	if err := zip.Close(); err != nil {
		return nil, err
	}

	return buf, nil
}

// Clean removes build artifacts from compiled runtimes.
func (f *Function) Clean() error {
	if err := f.RunHook("clean"); err != nil {
		return err
	}

	if r, ok := f.runtime.(runtime.CompiledRuntime); ok {
		return r.Clean(f.Path)
	}

	return nil
}

// A HookError represents a failed hook command.
type HookError struct {
	Name    string
	Command string
	Output  string
}

// Error string.
func (e *HookError) Error() string {
	return fmt.Sprintf("hook %q: %s", e.Name, e.Output)
}

// RunHook executes hook `name` in the function's directory.
func (f *Function) RunHook(name string) error {
	var command string

	switch name {
	case "clean":
		command = f.Hooks.Clean
	case "build":
		command = f.Hooks.Build
	}

	if command == "" {
		return nil
	}

	f.Log.WithFields(log.Fields{
		"hook":    name,
		"command": command,
	}).Debug("hook")

	cmd := exec.Command("sh", "-c", command)
	cmd.Env = os.Environ()
	cmd.Env = append(cmd.Env, fmt.Sprintf("FUNCTION=%s", f.Name))
	cmd.Dir = f.Path

	b, err := cmd.CombinedOutput()
	if err != nil {
		return &HookError{
			Name:    name,
			Command: command,
			Output:  string(b),
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package search

import (
	"container/heap"

	gr "github.com/gonum/graph"
)

// Sets up the cost functions and successor functions so I don't have to do a type switch every time.
// This almost always does more work than is necessary, but since it's only executed once per function, and graph functions are rather costly, the "extra work"
// should be negligible.
func setupFuncs(graph gr.Graph, cost, heuristicCost func(gr.Node, gr.Node) float64) (successorsFunc, predecessorsFunc, neighborsFunc func(gr.Node) []gr.Node, isSuccessorFunc, isPredecessorFunc,
	isNeighborFunc func(gr.Node, gr.Node) bool,
	costFunc, heuristicCostFunc func(gr.Node, gr.Node) float64) {

	switch g := graph.(type) {
	case gr.DirectedGraph:
		successorsFunc = g.Successors
		predecessorsFunc = g.Predecessors
		neighborsFunc = g.Neighbors
		isSuccessorFunc = g.IsSuccessor
		isPredecessorFunc = g.IsPredecessor
		isNeighborFunc = g.IsNeighbor
	default:
		successorsFunc = g.Neighbors
		predecessorsFunc = g.Neighbors
		neighborsFunc = g.Neighbors
		isSuccessorFunc = g.IsNeighbor
		isPredecessorFunc = g.IsNeighbor
		isNeighborFunc = g.IsNeighbor
	}

	if heuristicCost != nil {
		heuristicCostFunc = heuristicCost
	} else {
		if g, ok := graph.(gr.HeuristicCoster); ok {
			heuristicCostFunc = g.HeuristicCost
		} else {
			heuristicCostFunc = NullHeuristic
		}
	}

	if cost != nil {
		costFunc = cost
	} else {
		if g, ok := graph.(gr.Coster); ok {
			costFunc = g.Cost
		} else {
			costFunc = UniformCost
		}
	}

	return
}

/* Purely internal data structures and functions (mostly for sorting) */

// A package that contains an edge (as from EdgeList), and a Weight (as if Cost(Edge.Head(), Edge.Tail()) had been called)
type WeightedEdge struct {
	gr.Edge
	Weight float64
}

/** Sorts a list of edges by weight, agnostic to repeated edges as well as direction **/

type edgeSorter []WeightedEdge

func (el edgeSorter) Len() int {
	return len(el)
}

func (el edgeSorter) Less(i, j int) bool {
	return el[i].Weight < el[j].Weight
}

func (el edgeSorter) Swap(i, j int) {
	el[i], el[j] = el[j], el[i]
}

/** Keeps track of a node's scores so they can be used in a priority queue for A* **/

type internalNode struct {
	gr.Node
	gscore, fscore float64
}

/* A* stuff */
type aStarPriorityQueue struct {
	indexList map[int]int
	nodes     []internalNode
}

func (pq *aStarPriorityQueue) Less(i, j int) bool {
	return pq.nodes[i].fscore < pq.nodes[j].fscore // As the heap documentation says, a priority queue is listed if the actual values are treated as if they were negative
}

func (pq *aStarPriorityQueue) Swap(i, j int) {
	pq.indexList[pq.nodes[i].ID()] = j
	pq.indexList[pq.nodes[j].ID()] = i

	pq.nodes[i], pq.nodes[j] = pq.nodes[j], pq.nodes[i]
}

func (pq *aStarPriorityQueue) Len() int {
	return len(pq.nodes)
}

func (pq *aStarPriorityQueue) Push(x interface{}) {
	node := x.(internalNode)
	pq.nodes = append(pq.nodes, node)
	pq.indexList[node.ID()] = len(pq.nodes) - 1
}

func (pq *aStarPriorityQueue) Pop() interface{} {
	x := pq.nodes[len(pq.nodes)-1]
	pq.nodes = pq.nodes[:len(pq.nodes)-1]
	delete(pq.indexList, x.ID())

	return x
}

func (pq *aStarPriorityQueue) Fix(id int, newGScore, newFScore float64) {
	if i, ok := pq.indexList[id]; ok {
		pq.nodes[i].gscore = newGScore
		pq.nodes[i].fscore = newFScore
		heap.Fix(pq, i)
	}
}

func (pq *aStarPriorityQueue) Find(id int) (internalNode, bool) {
	loc, ok := pq.indexList[id]
	if ok {
		return pq.nodes[loc], true
	} else {
		return internalNode{}, false
	}

}

func (pq *aStarPriorityQueue) Exists(id int) bool {
	_, ok := pq.indexList[id]
	return ok
}

// General utility funcs

// Rebuilds a path backwards from the goal.
func rebuildPath(predecessors map[int]gr.Node, goal gr.Node) []gr.Node {
	if n, ok := goal.(internalNode); ok {
		goal = n.Node
	}
	path := []gr.Node{goal}
	curr := goal
	for prev, ok := predecessors[curr.ID()]; ok; prev, ok = predecessors[curr.ID()] {
		if n, ok := prev.(internalNode); ok {
			prev = n.Node
		}
		path = append(path, prev)
		curr = prev
	}

	// Reverse the path since it was built backwards
	for i, j := 0, len(path)-1; i < j; i, j = i+1, j-1 {
		path[i], path[j] = path[j], path[i]
	}

	return path
}
<commit_msg>Added way to sort nodes<commit_after>package search

import (
	"container/heap"

	gr "github.com/gonum/graph"
)

// Sets up the cost functions and successor functions so I don't have to do a type switch every time.
// This almost always does more work than is necessary, but since it's only executed once per function, and graph functions are rather costly, the "extra work"
// should be negligible.
func setupFuncs(graph gr.Graph, cost, heuristicCost func(gr.Node, gr.Node) float64) (successorsFunc, predecessorsFunc, neighborsFunc func(gr.Node) []gr.Node, isSuccessorFunc, isPredecessorFunc,
	isNeighborFunc func(gr.Node, gr.Node) bool,
	costFunc, heuristicCostFunc func(gr.Node, gr.Node) float64) {

	switch g := graph.(type) {
	case gr.DirectedGraph:
		successorsFunc = g.Successors
		predecessorsFunc = g.Predecessors
		neighborsFunc = g.Neighbors
		isSuccessorFunc = g.IsSuccessor
		isPredecessorFunc = g.IsPredecessor
		isNeighborFunc = g.IsNeighbor
	default:
		successorsFunc = g.Neighbors
		predecessorsFunc = g.Neighbors
		neighborsFunc = g.Neighbors
		isSuccessorFunc = g.IsNeighbor
		isPredecessorFunc = g.IsNeighbor
		isNeighborFunc = g.IsNeighbor
	}

	if heuristicCost != nil {
		heuristicCostFunc = heuristicCost
	} else {
		if g, ok := graph.(gr.HeuristicCoster); ok {
			heuristicCostFunc = g.HeuristicCost
		} else {
			heuristicCostFunc = NullHeuristic
		}
	}

	if cost != nil {
		costFunc = cost
	} else {
		if g, ok := graph.(gr.Coster); ok {
			costFunc = g.Cost
		} else {
			costFunc = UniformCost
		}
	}

	return
}

/* Purely internal data structures and functions (mostly for sorting) */

// A package that contains an edge (as from EdgeList), and a Weight (as if Cost(Edge.Head(), Edge.Tail()) had been called)
type WeightedEdge struct {
	gr.Edge
	Weight float64
}

/** Sorts a list of edges by weight, agnostic to repeated edges as well as direction **/

type edgeSorter []WeightedEdge

func (el edgeSorter) Len() int {
	return len(el)
}

func (el edgeSorter) Less(i, j int) bool {
	return el[i].Weight < el[j].Weight
}

func (el edgeSorter) Swap(i, j int) {
	el[i], el[j] = el[j], el[i]
}

/** Keeps track of a node's scores so they can be used in a priority queue for A* **/

type internalNode struct {
	gr.Node
	gscore, fscore float64
}

/* A* stuff */
type aStarPriorityQueue struct {
	indexList map[int]int
	nodes     []internalNode
}

func (pq *aStarPriorityQueue) Less(i, j int) bool {
	return pq.nodes[i].fscore < pq.nodes[j].fscore // As the heap documentation says, a priority queue is listed if the actual values are treated as if they were negative
}

func (pq *aStarPriorityQueue) Swap(i, j int) {
	pq.indexList[pq.nodes[i].ID()] = j
	pq.indexList[pq.nodes[j].ID()] = i

	pq.nodes[i], pq.nodes[j] = pq.nodes[j], pq.nodes[i]
}

func (pq *aStarPriorityQueue) Len() int {
	return len(pq.nodes)
}

func (pq *aStarPriorityQueue) Push(x interface{}) {
	node := x.(internalNode)
	pq.nodes = append(pq.nodes, node)
	pq.indexList[node.ID()] = len(pq.nodes) - 1
}

func (pq *aStarPriorityQueue) Pop() interface{} {
	x := pq.nodes[len(pq.nodes)-1]
	pq.nodes = pq.nodes[:len(pq.nodes)-1]
	delete(pq.indexList, x.ID())

	return x
}

func (pq *aStarPriorityQueue) Fix(id int, newGScore, newFScore float64) {
	if i, ok := pq.indexList[id]; ok {
		pq.nodes[i].gscore = newGScore
		pq.nodes[i].fscore = newFScore
		heap.Fix(pq, i)
	}
}

func (pq *aStarPriorityQueue) Find(id int) (internalNode, bool) {
	loc, ok := pq.indexList[id]
	if ok {
		return pq.nodes[loc], true
	} else {
		return internalNode{}, false
	}

}

func (pq *aStarPriorityQueue) Exists(id int) bool {
	_, ok := pq.indexList[id]
	return ok
}

type denseNodeSorter []gr.Node

func (dns denseNodeSorter) Less(i, j int) bool {
	return dns[i].ID() < dns[j].ID()
}

func (dns denseNodeSorter) Swap(i, j int) {
	dns[i], dns[j] = dns[j], dns[i]
}

func (dns denseNodeSorter) Len() int {
	return len(dns)
}

// General utility funcs

// Rebuilds a path backwards from the goal.
func rebuildPath(predecessors map[int]gr.Node, goal gr.Node) []gr.Node {
	if n, ok := goal.(internalNode); ok {
		goal = n.Node
	}
	path := []gr.Node{goal}
	curr := goal
	for prev, ok := predecessors[curr.ID()]; ok; prev, ok = predecessors[curr.ID()] {
		if n, ok := prev.(internalNode); ok {
			prev = n.Node
		}
		path = append(path, prev)
		curr = prev
	}

	// Reverse the path since it was built backwards
	for i, j := 0, len(path)-1; i < j; i, j = i+1, j-1 {
		path[i], path[j] = path[j], path[i]
	}

	return path
}
<|endoftext|>