<commit_before>package collector

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/karimra/gnmic/outputs"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/encoding/prototext"
)

const (
	defaultTargetReceivebuffer = 1000
)

// Config is the collector config
type Config struct {
	PrometheusAddress   string
	Debug               bool
	Format              string
	TargetReceiveBuffer uint
}

// Collector //
type Collector struct {
	Config        *Config
	Subscriptions map[string]*SubscriptionConfig
	Outputs       map[string][]outputs.Output
	DialOpts      []grpc.DialOption
	//
	m          *sync.Mutex
	Targets    map[string]*Target
	Logger     *log.Logger
	httpServer *http.Server

	ctx      context.Context
	cancelFn context.CancelFunc
}

// NewCollector //
func NewCollector(ctx context.Context,
	config *Config,
	targetConfigs map[string]*TargetConfig,
	subscriptions map[string]*SubscriptionConfig,
	outputs map[string][]outputs.Output,
	dialOpts []grpc.DialOption,
	logger *log.Logger,
) *Collector {
	nctx, cancel := context.WithCancel(ctx)
	grpcMetrics := grpc_prometheus.NewClientMetrics()
	reg := prometheus.NewRegistry()
	reg.MustRegister(prometheus.NewGoCollector())
	grpcMetrics.EnableClientHandlingTimeHistogram()
	reg.MustRegister(grpcMetrics)
	handler := http.NewServeMux()
	handler.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
	httpServer := &http.Server{
		Handler: handler,
		Addr:    config.PrometheusAddress,
	}
	dialOpts = append(dialOpts, grpc.WithStreamInterceptor(grpcMetrics.StreamClientInterceptor()))
	if config.TargetReceiveBuffer == 0 {
		config.TargetReceiveBuffer = defaultTargetReceivebuffer
	}
	c := &Collector{
		Config:        config,
		Subscriptions: subscriptions,
		Outputs:       outputs,
		DialOpts:      dialOpts,
		m:             new(sync.Mutex),
		Targets:       make(map[string]*Target),
		Logger:        logger,
		httpServer:    httpServer,
		ctx:           nctx,
		cancelFn:      cancel,
	}
	wg := new(sync.WaitGroup)
	wg.Add(len(targetConfigs))
	for _, tc := range targetConfigs {
		go func(tc *TargetConfig) {
			defer wg.Done()
			err := c.InitTarget(tc)
			if err == context.DeadlineExceeded {
				c.Logger.Printf("failed to initialize target '%s' timeout (%s) reached: %v", tc.Name, tc.Timeout, err)
				return
			}
			if err != nil {
				c.Logger.Printf("failed to initialize target '%s': %v", tc.Name, err)
				return
			}
			c.Logger.Printf("target '%s' initialized", tc.Name)
		}(tc)
	}
	wg.Wait()
	return c
}

// InitTarget initializes a target based on *TargetConfig
func (c *Collector) InitTarget(tc *TargetConfig) error {
	if tc.BufferSize == 0 {
		tc.BufferSize = c.Config.TargetReceiveBuffer
	}
	t := NewTarget(tc)
	//
	t.Subscriptions = make([]*SubscriptionConfig, 0, len(tc.Subscriptions))
	for _, subName := range tc.Subscriptions {
		if sub, ok := c.Subscriptions[subName]; ok {
			t.Subscriptions = append(t.Subscriptions, sub)
		}
	}
	if len(t.Subscriptions) == 0 {
		t.Subscriptions = make([]*SubscriptionConfig, 0, len(c.Subscriptions))
		for _, sub := range c.Subscriptions {
			t.Subscriptions = append(t.Subscriptions, sub)
		}
	}
	//
	t.Outputs = make([]outputs.Output, 0, len(tc.Outputs))
	for _, outName := range tc.Outputs {
		if outs, ok := c.Outputs[outName]; ok {
			t.Outputs = append(t.Outputs, outs...)
		}
	}
	if len(t.Outputs) == 0 {
		t.Outputs = make([]outputs.Output, 0, len(c.Outputs))
		for _, o := range c.Outputs {
			t.Outputs = append(t.Outputs, o...)
		}
	}
	//
	err := t.CreateGNMIClient(c.ctx, c.DialOpts...)
	if err != nil {
		return err
	}
	c.m.Lock()
	defer c.m.Unlock()
	c.Targets[t.Config.Name] = t
	return nil
}

// Subscribe //
func (c *Collector) Subscribe(tName string) error {
	if t, ok := c.Targets[tName]; ok {
		for _, sc := range t.Subscriptions {
			req, err := sc.CreateSubscribeRequest()
			if err != nil {
				return err
			}
			c.Logger.Printf("sending gNMI SubscribeRequest: subscribe='%+v', mode='%+v', encoding='%+v', to %s",
				req, req.GetSubscribe().GetMode(), req.GetSubscribe().GetEncoding(), t.Config.Name)
			go t.Subscribe(c.ctx, req, sc.Name)
		}
		return nil
	}
	return fmt.Errorf("unknown target name: %s", tName)
}

// Start start the prometheus server as well as a goroutine per target selecting on the response chan, the error chan and the ctx.Done() chan
func (c *Collector) Start() {
	go func() {
		if err := c.httpServer.ListenAndServe(); err != nil {
			c.Logger.Printf("Unable to start prometheus http server: %v", err)
			return
		}
	}()
	defer func() {
		for _, outputs := range c.Outputs {
			for _, o := range outputs {
				o.Close()
			}
		}
	}()
	wg := new(sync.WaitGroup)
	wg.Add(len(c.Targets))
	for _, t := range c.Targets {
		go func(t *Target) {
			defer wg.Done()
			numOnceSubscriptions := t.numberOfOnceSubscriptions()
			remainingOnceSubscriptions := numOnceSubscriptions
			numSubscriptions := len(t.Subscriptions)
			for {
				select {
				case rsp := <-t.SubscribeResponses:
					c.Logger.Printf("received subscribe response: %+v", rsp)
					m := make(map[string]interface{})
					m["subscription-name"] = rsp.SubscriptionName
					m["source"] = t.Config.Name
					b, err := c.FormatMsg(m, rsp.Response)
					if err != nil {
						c.Logger.Printf("failed formatting msg from target '%s': %v", t.Config.Name, err)
						continue
					}
					if c.subscriptionMode(rsp.SubscriptionName) == "ONCE" {
						t.Export(b, outputs.Meta{"source": t.Config.Name, "format": c.Config.Format})
					} else {
						go t.Export(b, outputs.Meta{"source": t.Config.Name, "format": c.Config.Format})
					}
					if remainingOnceSubscriptions > 0 {
						if c.subscriptionMode(rsp.SubscriptionName) == "ONCE" {
							switch rsp.Response.Response.(type) {
							case *gnmi.SubscribeResponse_SyncResponse:
								remainingOnceSubscriptions--
							}
						}
					}
					if remainingOnceSubscriptions == 0 && numSubscriptions == numOnceSubscriptions {
						return
					}
				case err := <-t.Errors:
					if err == io.EOF {
						c.Logger.Printf("target '%s' closed stream(EOF)", t.Config.Name)
						return
					}
					c.Logger.Printf("target '%s' error: %v", t.Config.Name, err)
				}
			}
		}(t)
	}
	wg.Wait()
}

// FormatMsg formats the gnmi.SubscribeResponse and returns a []byte and an error
func (c *Collector) FormatMsg(meta map[string]interface{}, rsp *gnmi.SubscribeResponse) ([]byte, error) {
	if rsp == nil {
		return nil, nil
	}
	if c.Config.Format == "textproto" {
		return []byte(prototext.Format(rsp)), nil
	}
	switch rsp := rsp.Response.(type) {
	case *gnmi.SubscribeResponse_Update:
		msg := new(msg)
		msg.Timestamp = rsp.Update.Timestamp
		t := time.Unix(0, rsp.Update.Timestamp)
		msg.Time = &t
		if meta == nil {
			meta = make(map[string]interface{})
		}
		msg.Prefix = gnmiPathToXPath(rsp.Update.Prefix)
		var ok bool
		if _, ok = meta["source"]; ok {
			msg.Source = fmt.Sprintf("%s", meta["source"])
		}
		if _, ok = meta["system-name"]; ok {
			msg.SystemName = fmt.Sprintf("%s", meta["system-name"])
		}
		if _, ok = meta["subscription-name"]; ok {
			msg.SubscriptionName = fmt.Sprintf("%s", meta["subscription-name"])
		}
		for i, upd := range rsp.Update.Update {
			pathElems := make([]string, 0, len(upd.Path.Elem))
			for _, pElem := range upd.Path.Elem {
				pathElems = append(pathElems, pElem.GetName())
			}
			value, err := getValue(upd.Val)
			if err != nil {
				c.Logger.Println(err)
			}
			msg.Updates = append(msg.Updates,
				&update{
					Path:   gnmiPathToXPath(upd.Path),
					Values: make(map[string]interface{}),
				})
			msg.Updates[i].Values[strings.Join(pathElems, "/")] = value
		}
		for _, del := range rsp.Update.Delete {
			msg.Deletes = append(msg.Deletes, gnmiPathToXPath(del))
		}
		data, err := json.Marshal(msg)
		if err != nil {
			return nil, err
		}
		return data, nil
	}
	return nil, nil
}

// TargetPoll sends a gnmi.SubscribeRequest_Poll to targetName and returns the response and an error,
// it uses the targetName and the subscriptionName strings to find the gnmi.GNMI_SubscribeClient
func (c *Collector) TargetPoll(targetName, subscriptionName string) (*gnmi.SubscribeResponse, error) {
	if sub, ok := c.Subscriptions[subscriptionName]; ok {
		if strings.ToUpper(sub.Mode) != "POLL" {
			return nil, fmt.Errorf("subscription '%s' is not a POLL subscription", subscriptionName)
		}
		if t, ok := c.Targets[targetName]; ok {
			if subClient, ok := t.SubscribeClients[subscriptionName]; ok {
				err := subClient.Send(&gnmi.SubscribeRequest{
					Request: &gnmi.SubscribeRequest_Poll{
						Poll: &gnmi.Poll{},
					},
				})
				if err != nil {
					return nil, err
				}
				return subClient.Recv()
			}
		}
		return nil, fmt.Errorf("unknown target name '%s'", targetName)
	}
	return nil, fmt.Errorf("unknown subscription name '%s'", subscriptionName)
}

// PolledSubscriptionsTargets returns a map of target name to a list of subscription names that have Mode == POLL
func (c *Collector) PolledSubscriptionsTargets() map[string][]string {
	result := make(map[string][]string)
	for tn, target := range c.Targets {
		for _, sub := range target.Subscriptions {
			if strings.ToUpper(sub.Mode) == "POLL" {
				if result[tn] == nil {
					result[tn] = make([]string, 0)
				}
				result[tn] = append(result[tn], sub.Name)
			}
		}
	}
	return result
}

func (c *Collector) subscriptionMode(name string) string {
	if sub, ok := c.Subscriptions[name]; ok {
		return strings.ToUpper(sub.Mode)
	}
	return ""
}
<commit_msg>move subscribe response logging under debug<commit_after>package collector

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/karimra/gnmic/outputs"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/encoding/prototext"
)

const (
	defaultTargetReceivebuffer = 1000
)

// Config is the collector config
type Config struct {
	PrometheusAddress   string
	Debug               bool
	Format              string
	TargetReceiveBuffer uint
}

// Collector //
type Collector struct {
	Config        *Config
	Subscriptions map[string]*SubscriptionConfig
	Outputs       map[string][]outputs.Output
	DialOpts      []grpc.DialOption
	//
	m          *sync.Mutex
	Targets    map[string]*Target
	Logger     *log.Logger
	httpServer *http.Server

	ctx      context.Context
	cancelFn context.CancelFunc
}

// NewCollector //
func NewCollector(ctx context.Context,
	config *Config,
	targetConfigs map[string]*TargetConfig,
	subscriptions map[string]*SubscriptionConfig,
	outputs map[string][]outputs.Output,
	dialOpts []grpc.DialOption,
	logger *log.Logger,
) *Collector {
	nctx, cancel := context.WithCancel(ctx)
	grpcMetrics := grpc_prometheus.NewClientMetrics()
	reg := prometheus.NewRegistry()
	reg.MustRegister(prometheus.NewGoCollector())
	grpcMetrics.EnableClientHandlingTimeHistogram()
	reg.MustRegister(grpcMetrics)
	handler := http.NewServeMux()
	handler.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
	httpServer := &http.Server{
		Handler: handler,
		Addr:    config.PrometheusAddress,
	}
	dialOpts = append(dialOpts, grpc.WithStreamInterceptor(grpcMetrics.StreamClientInterceptor()))
	if config.TargetReceiveBuffer == 0 {
		config.TargetReceiveBuffer = defaultTargetReceivebuffer
	}
	c := &Collector{
		Config:        config,
		Subscriptions: subscriptions,
		Outputs:       outputs,
		DialOpts:      dialOpts,
		m:             new(sync.Mutex),
		Targets:       make(map[string]*Target),
		Logger:        logger,
		httpServer:    httpServer,
		ctx:           nctx,
		cancelFn:      cancel,
	}
	wg := new(sync.WaitGroup)
	wg.Add(len(targetConfigs))
	for _, tc := range targetConfigs {
		go func(tc *TargetConfig) {
			defer wg.Done()
			err := c.InitTarget(tc)
			if err == context.DeadlineExceeded {
				c.Logger.Printf("failed to initialize target '%s' timeout (%s) reached: %v", tc.Name, tc.Timeout, err)
				return
			}
			if err != nil {
				c.Logger.Printf("failed to initialize target '%s': %v", tc.Name, err)
				return
			}
			c.Logger.Printf("target '%s' initialized", tc.Name)
		}(tc)
	}
	wg.Wait()
	return c
}

// InitTarget initializes a target based on *TargetConfig
func (c *Collector) InitTarget(tc *TargetConfig) error {
	if tc.BufferSize == 0 {
		tc.BufferSize = c.Config.TargetReceiveBuffer
	}
	t := NewTarget(tc)
	//
	t.Subscriptions = make([]*SubscriptionConfig, 0, len(tc.Subscriptions))
	for _, subName := range tc.Subscriptions {
		if sub, ok := c.Subscriptions[subName]; ok {
			t.Subscriptions = append(t.Subscriptions, sub)
		}
	}
	if len(t.Subscriptions) == 0 {
		t.Subscriptions = make([]*SubscriptionConfig, 0, len(c.Subscriptions))
		for _, sub := range c.Subscriptions {
			t.Subscriptions = append(t.Subscriptions, sub)
		}
	}
	//
	t.Outputs = make([]outputs.Output, 0, len(tc.Outputs))
	for _, outName := range tc.Outputs {
		if outs, ok := c.Outputs[outName]; ok {
			t.Outputs = append(t.Outputs, outs...)
		}
	}
	if len(t.Outputs) == 0 {
		t.Outputs = make([]outputs.Output, 0, len(c.Outputs))
		for _, o := range c.Outputs {
			t.Outputs = append(t.Outputs, o...)
		}
	}
	//
	err := t.CreateGNMIClient(c.ctx, c.DialOpts...)
	if err != nil {
		return err
	}
	c.m.Lock()
	defer c.m.Unlock()
	c.Targets[t.Config.Name] = t
	return nil
}

// Subscribe //
func (c *Collector) Subscribe(tName string) error {
	if t, ok := c.Targets[tName]; ok {
		for _, sc := range t.Subscriptions {
			req, err := sc.CreateSubscribeRequest()
			if err != nil {
				return err
			}
			c.Logger.Printf("sending gNMI SubscribeRequest: subscribe='%+v', mode='%+v', encoding='%+v', to %s",
				req, req.GetSubscribe().GetMode(), req.GetSubscribe().GetEncoding(), t.Config.Name)
			go t.Subscribe(c.ctx, req, sc.Name)
		}
		return nil
	}
	return fmt.Errorf("unknown target name: %s", tName)
}

// Start start the prometheus server as well as a goroutine per target selecting on the response chan, the error chan and the ctx.Done() chan
func (c *Collector) Start() {
	go func() {
		if err := c.httpServer.ListenAndServe(); err != nil {
			c.Logger.Printf("Unable to start prometheus http server: %v", err)
			return
		}
	}()
	defer func() {
		for _, outputs := range c.Outputs {
			for _, o := range outputs {
				o.Close()
			}
		}
	}()
	wg := new(sync.WaitGroup)
	wg.Add(len(c.Targets))
	for _, t := range c.Targets {
		go func(t *Target) {
			defer wg.Done()
			numOnceSubscriptions := t.numberOfOnceSubscriptions()
			remainingOnceSubscriptions := numOnceSubscriptions
			numSubscriptions := len(t.Subscriptions)
			for {
				select {
				case rsp := <-t.SubscribeResponses:
					if c.Config.Debug {
						c.Logger.Printf("received subscribe response: %+v", rsp)
					}
					m := make(map[string]interface{})
					m["subscription-name"] = rsp.SubscriptionName
					m["source"] = t.Config.Name
					b, err := c.FormatMsg(m, rsp.Response)
					if err != nil {
						c.Logger.Printf("failed formatting msg from target '%s': %v", t.Config.Name, err)
						continue
					}
					if c.subscriptionMode(rsp.SubscriptionName) == "ONCE" {
						t.Export(b, outputs.Meta{"source": t.Config.Name, "format": c.Config.Format})
					} else {
						go t.Export(b, outputs.Meta{"source": t.Config.Name, "format": c.Config.Format})
					}
					if remainingOnceSubscriptions > 0 {
						if c.subscriptionMode(rsp.SubscriptionName) == "ONCE" {
							switch rsp.Response.Response.(type) {
							case *gnmi.SubscribeResponse_SyncResponse:
								remainingOnceSubscriptions--
							}
						}
					}
					if remainingOnceSubscriptions == 0 && numSubscriptions == numOnceSubscriptions {
						return
					}
				case err := <-t.Errors:
					if err == io.EOF {
						c.Logger.Printf("target '%s' closed stream(EOF)", t.Config.Name)
						return
					}
					c.Logger.Printf("target '%s' error: %v", t.Config.Name, err)
				}
			}
		}(t)
	}
	wg.Wait()
}

// FormatMsg formats the gnmi.SubscribeResponse and returns a []byte and an error
func (c *Collector) FormatMsg(meta map[string]interface{}, rsp *gnmi.SubscribeResponse) ([]byte, error) {
	if rsp == nil {
		return nil, nil
	}
	if c.Config.Format == "textproto" {
		return []byte(prototext.Format(rsp)), nil
	}
	switch rsp := rsp.Response.(type) {
	case *gnmi.SubscribeResponse_Update:
		msg := new(msg)
		msg.Timestamp = rsp.Update.Timestamp
		t := time.Unix(0, rsp.Update.Timestamp)
		msg.Time = &t
		if meta == nil {
			meta = make(map[string]interface{})
		}
		msg.Prefix = gnmiPathToXPath(rsp.Update.Prefix)
		var ok bool
		if _, ok = meta["source"]; ok {
			msg.Source = fmt.Sprintf("%s", meta["source"])
		}
		if _, ok = meta["system-name"]; ok {
			msg.SystemName = fmt.Sprintf("%s", meta["system-name"])
		}
		if _, ok = meta["subscription-name"]; ok {
			msg.SubscriptionName = fmt.Sprintf("%s", meta["subscription-name"])
		}
		for i, upd := range rsp.Update.Update {
			pathElems := make([]string, 0, len(upd.Path.Elem))
			for _, pElem := range upd.Path.Elem {
				pathElems = append(pathElems, pElem.GetName())
			}
			value, err := getValue(upd.Val)
			if err != nil {
				c.Logger.Println(err)
			}
			msg.Updates = append(msg.Updates,
				&update{
					Path:   gnmiPathToXPath(upd.Path),
					Values: make(map[string]interface{}),
				})
			msg.Updates[i].Values[strings.Join(pathElems, "/")] = value
		}
		for _, del := range rsp.Update.Delete {
			msg.Deletes = append(msg.Deletes, gnmiPathToXPath(del))
		}
		data, err := json.Marshal(msg)
		if err != nil {
			return nil, err
		}
		return data, nil
	}
	return nil, nil
}

// TargetPoll sends a gnmi.SubscribeRequest_Poll to targetName and returns the response and an error,
// it uses the targetName and the subscriptionName strings to find the gnmi.GNMI_SubscribeClient
func (c *Collector) TargetPoll(targetName, subscriptionName string) (*gnmi.SubscribeResponse, error) {
	if sub, ok := c.Subscriptions[subscriptionName]; ok {
		if strings.ToUpper(sub.Mode) != "POLL" {
			return nil, fmt.Errorf("subscription '%s' is not a POLL subscription", subscriptionName)
		}
		if t, ok := c.Targets[targetName]; ok {
			if subClient, ok := t.SubscribeClients[subscriptionName]; ok {
				err := subClient.Send(&gnmi.SubscribeRequest{
					Request: &gnmi.SubscribeRequest_Poll{
						Poll: &gnmi.Poll{},
					},
				})
				if err != nil {
					return nil, err
				}
				return subClient.Recv()
			}
		}
		return nil, fmt.Errorf("unknown target name '%s'", targetName)
	}
	return nil, fmt.Errorf("unknown subscription name '%s'", subscriptionName)
}

// PolledSubscriptionsTargets returns a map of target name to a list of subscription names that have Mode == POLL
func (c *Collector) PolledSubscriptionsTargets() map[string][]string {
	result := make(map[string][]string)
	for tn, target := range c.Targets {
		for _, sub := range target.Subscriptions {
			if strings.ToUpper(sub.Mode) == "POLL" {
				if result[tn] == nil {
					result[tn] = make([]string, 0)
				}
				result[tn] = append(result[tn], sub.Name)
			}
		}
	}
	return result
}

func (c *Collector) subscriptionMode(name string) string {
	if sub, ok := c.Subscriptions[name]; ok {
		return strings.ToUpper(sub.Mode)
	}
	return ""
}
<|endoftext|>


<commit_before>package collector

import (
	"launchpad.net/goyaml"
	"database/sql"
	_ "github.com/mattn/go-sqlite3"
	"fmt"
	"os/exec"
)

type Collector struct{}

type Unit struct {
	State string
}

type Service struct {
	Units map[string]Unit
}

type output struct {
	Services map[string]Service
}

func (c *Collector) Collect() ([]byte, error) {
	fmt.Println("collecting status")
	return exec.Command("juju", "status").Output()
}

func (c *Collector) Parse(data []byte) *output {
	fmt.Println("parsing yaml")
	raw := new(output)
	_ = goyaml.Unmarshal(data, raw)
	return raw
}

func (c *Collector) Update(db *sql.DB, out *output) {
	fmt.Println("updating status")

	updateApp, _ := db.Prepare("UPDATE apps SET state=?")

	for _, service := range out.Services {
		for _, unit := range service.Units {
			tx, _ := db.Begin()
			stmt := tx.Stmt(updateApp)
			defer stmt.Close()
			if unit.State == "started" {
				stmt.Exec("STARTED")
			} else {
				stmt.Exec("STOPPED")
			}
			tx.Commit()
		}
	}

}
<commit_msg>fixed update to set state app by app<commit_after>package collector

import (
	"launchpad.net/goyaml"
	"database/sql"
	_ "github.com/mattn/go-sqlite3"
	"fmt"
	"os/exec"
)

type Collector struct{}

type Unit struct {
	State string
}

type Service struct {
	Units map[string]Unit
}

type output struct {
	Services map[string]Service
}

func (c *Collector) Collect() ([]byte, error) {
	fmt.Println("collecting status")
	return exec.Command("juju", "status").Output()
}

func (c *Collector) Parse(data []byte) *output {
	fmt.Println("parsing yaml")
	raw := new(output)
	_ = goyaml.Unmarshal(data, raw)
	return raw
}

func (c *Collector) Update(db *sql.DB, out *output) {
	fmt.Println("updating status")

	var state string

	updateApp, _ := db.Prepare("UPDATE apps SET state=? WHERE name=?")

	for serviceName, service := range out.Services {
		for _, unit := range service.Units {
			tx, _ := db.Begin()
			stmt := tx.Stmt(updateApp)
			defer stmt.Close()
			if unit.State == "started" {
				state = "STARTED"
			} else {
				state = "STOPPED"
			}
			stmt.Exec(state, serviceName)
			tx.Commit()
		}
	}

}
<|endoftext|>


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

package inode_test

import (
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/googlecloudplatform/gcsfuse/fs/inode"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	. "github.com/jacobsa/ogletest"
)

func TestDir(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

const inodeID = 17
const inodeName = "foo/bar/"
const typeCacheTTL = time.Second

type DirTest struct {
	ctx    context.Context
	bucket gcs.Bucket
	clock  timeutil.SimulatedClock

	in *inode.DirInode
}

var _ SetUpInterface = &DirTest{}
var _ TearDownInterface = &DirTest{}

func init() { RegisterTestSuite(&DirTest{}) }

func (t *DirTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
	t.clock.SetTime(time.Date(2015, 4, 5, 2, 15, 0, 0, time.Local))
	t.bucket = gcsfake.NewFakeBucket(&t.clock, "some_bucket")

	// Create the inode. No implicit dirs by default.
	t.resetInode(false)
}

func (t *DirTest) TearDown() {
	t.in.Unlock()
}

func (t *DirTest) resetInode(implicitDirs bool) {
	if t.in != nil {
		t.in.Unlock()
	}

	t.in = inode.NewDirInode(
		inodeID,
		inodeName,
		implicitDirs,
		typeCacheTTL,
		t.bucket,
		&t.clock)

	t.in.Lock()
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *DirTest) ID() {
	ExpectEq(inodeID, t.in.ID())
}

func (t *DirTest) Name() {
	ExpectEq(inodeName, t.in.Name())
}

func (t *DirTest) LookupCount() {
	// Increment thrice. The count should now be three.
	t.in.IncrementLookupCount()
	t.in.IncrementLookupCount()
	t.in.IncrementLookupCount()

	// Decrementing twice shouldn't cause destruction. But one more should.
	AssertFalse(t.in.DecrementLookupCount(2))
	ExpectTrue(t.in.DecrementLookupCount(1))
}

func (t *DirTest) Attributes() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_NonExistent() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_FileOnly() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_DirOnly() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_ImplicitDirOnly_Disabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_ImplicitDirOnly_Enabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_FileAndDir() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_FileAndDirAndImplicitDir_Disabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_FileAndDirAndImplicitDir_Enabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_TypeCaching() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) ReadEntries_Empty() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) ReadEntries_NonEmpty_ImplicitDirsDisabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) ReadEntries_NonEmpty_ImplicitDirsEnabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) ReadEntries_LotsOfEntries() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) ReadEntries_TypeCaching() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildFile_DoesntExist() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildFile_Exists() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildFile_TypeCaching() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildDir_DoesntExist() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildDir_Exists() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildDir_TypeCaching() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) DeleteChildFile_DoesntExist() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) DeleteChildFile_Exists() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) DeleteChildFile_TypeCaching() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) DeleteChildDir_DoesntExist() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) DeleteChildDir_Exists() {
	AssertTrue(false, "TODO")
}

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

package inode_test

import (
	"os"
	"testing"
	"time"

	"golang.org/x/net/context"

	"github.com/googlecloudplatform/gcsfuse/fs/inode"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcsfake"
	. "github.com/jacobsa/ogletest"
)

func TestDir(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

const inodeID = 17
const inodeName = "foo/bar/"
const typeCacheTTL = time.Second

type DirTest struct {
	ctx    context.Context
	bucket gcs.Bucket
	clock  timeutil.SimulatedClock

	in *inode.DirInode
}

var _ SetUpInterface = &DirTest{}
var _ TearDownInterface = &DirTest{}

func init() { RegisterTestSuite(&DirTest{}) }

func (t *DirTest) SetUp(ti *TestInfo) {
	t.ctx = ti.Ctx
	t.clock.SetTime(time.Date(2015, 4, 5, 2, 15, 0, 0, time.Local))
	t.bucket = gcsfake.NewFakeBucket(&t.clock, "some_bucket")

	// Create the inode. No implicit dirs by default.
	t.resetInode(false)
}

func (t *DirTest) TearDown() {
	t.in.Unlock()
}

func (t *DirTest) resetInode(implicitDirs bool) {
	if t.in != nil {
		t.in.Unlock()
	}

	t.in = inode.NewDirInode(
		inodeID,
		inodeName,
		implicitDirs,
		typeCacheTTL,
		t.bucket,
		&t.clock)

	t.in.Lock()
}

////////////////////////////////////////////////////////////////////////
// Tests
////////////////////////////////////////////////////////////////////////

func (t *DirTest) ID() {
	ExpectEq(inodeID, t.in.ID())
}

func (t *DirTest) Name() {
	ExpectEq(inodeName, t.in.Name())
}

func (t *DirTest) LookupCount() {
	// Increment thrice. The count should now be three.
	t.in.IncrementLookupCount()
	t.in.IncrementLookupCount()
	t.in.IncrementLookupCount()

	// Decrementing twice shouldn't cause destruction. But one more should.
	AssertFalse(t.in.DecrementLookupCount(2))
	ExpectTrue(t.in.DecrementLookupCount(1))
}

func (t *DirTest) Attributes() {
	attrs, err := t.in.Attributes(context.Background())
	AssertEq(nil, err)
	ExpectEq(os.FileMode(0700)|os.ModeDir, attrs.Mode)
}

func (t *DirTest) LookUpChild_NonExistent() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_FileOnly() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_DirOnly() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_ImplicitDirOnly_Disabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_ImplicitDirOnly_Enabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_FileAndDir() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_FileAndDirAndImplicitDir_Disabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_FileAndDirAndImplicitDir_Enabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) LookUpChild_TypeCaching() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) ReadEntries_Empty() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) ReadEntries_NonEmpty_ImplicitDirsDisabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) ReadEntries_NonEmpty_ImplicitDirsEnabled() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) ReadEntries_LotsOfEntries() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) ReadEntries_TypeCaching() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildFile_DoesntExist() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildFile_Exists() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildFile_TypeCaching() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildDir_DoesntExist() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildDir_Exists() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) CreateChildDir_TypeCaching() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) DeleteChildFile_DoesntExist() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) DeleteChildFile_Exists() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) DeleteChildFile_TypeCaching() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) DeleteChildDir_DoesntExist() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) DeleteChildDir_Exists() {
	AssertTrue(false, "TODO")
}

func (t *DirTest) DeleteChildDir_TypeCaching() {
	AssertTrue(false, "TODO")
}
<|endoftext|>


<commit_before>// Copyright 2016 Joel Scoble and The JoeFriday 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 diskusage calculates IO usage of the block devices. Usage is
// calculated by taking the difference between two snapshots of IO statistics
// for block devices, /procd/diskstats. The time elapsed between the two
// snapshots is stored in the TimeDelta field.
package diskusage

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

	"github.com/SermoDigital/helpers"
	joe "github.com/mohae/joefriday"
	stats "github.com/mohae/joefriday/disk/diskstats"
	"github.com/mohae/joefriday/disk/structs"
)

// Profiler is used to process the IO usage of the block devices.
type Profiler struct {
	*stats.Profiler
	prior *structs.DiskStats
}

// Returns an initialized Profiler; ready to use. Upon creation, a
// /proc/diskstats snapshot is taken so that any Get() will return valid
// information.
func NewProfiler() (prof *Profiler, err error) {
	p, err := stats.NewProfiler()
	if err != nil {
		return nil, err
	}
	s, err := p.Get()
	if err != nil {
		return nil, err
	}
	return &Profiler{Profiler: p, prior: s}, nil
}

// Get returns the current IO usage of the block devices. Calculating usage
// requires two snapshots. This func gets the current snapshot of
// /proc/diskstats and calculates the difference between that and the prior
// snapshot. The current snapshot is stored for use as the prior snapshot on
// the next Get call. If ongoing usage information is desired, the Ticker
// should be used; it's better suited for ongoing usage information.
func (prof *Profiler) Get() (u *structs.DiskUsage, err error) {
	st, err := prof.Profiler.Get()
	if err != nil {
		return nil, err
	}
	u = prof.CalculateUsage(st)
	prof.prior = st
	return u, nil
}

var std *Profiler
var stdMu sync.Mutex

// Get returns the current IO usage of the block devices using the package's
// global Profiler. The Profiler is instantiated lazily. If the profiler
// doesn't already exist, the first usage information will not be useful due
// to minimal time elapsing between the initial and second snapshots used for
// usage calculations; the results of the first call should be discarded.
func Get() (u *structs.DiskUsage, err error) {
	stdMu.Lock()
	defer stdMu.Unlock()
	if std == nil {
		std, err = NewProfiler()
		if err != nil {
			return nil, err
		}
	}
	return std.Get()
}

// CalculateUsage returns the difference between the current /proc/diskstats
// snapshot and the prior one.
func (prof *Profiler) CalculateUsage(cur *structs.DiskStats) *structs.DiskUsage {
	u := &structs.DiskUsage{Timestamp: cur.Timestamp, Device: make([]structs.Device, len(cur.Device))}
	u.TimeDelta = cur.Timestamp - prof.prior.Timestamp
	for i := 0; i < len(cur.Device); i++ {
		u.Device[i].Major = cur.Device[i].Major
		u.Device[i].Minor = cur.Device[i].Minor
		u.Device[i].Name = cur.Device[i].Name
		u.Device[i].ReadsCompleted = cur.Device[i].ReadsCompleted - prof.prior.Device[i].ReadsCompleted
		u.Device[i].ReadsMerged = cur.Device[i].ReadsMerged - prof.prior.Device[i].ReadsMerged
		u.Device[i].ReadSectors = cur.Device[i].ReadSectors - prof.prior.Device[i].ReadSectors
		u.Device[i].ReadingTime = cur.Device[i].ReadingTime - prof.prior.Device[i].ReadingTime
		u.Device[i].WritesCompleted = cur.Device[i].WritesCompleted - prof.prior.Device[i].WritesCompleted
		u.Device[i].WritesMerged = cur.Device[i].WritesMerged - prof.prior.Device[i].WritesMerged
		u.Device[i].WrittenSectors = cur.Device[i].WrittenSectors - prof.prior.Device[i].WrittenSectors
		u.Device[i].WritingTime = cur.Device[i].WritingTime - prof.prior.Device[i].WritingTime
		u.Device[i].IOInProgress = cur.Device[i].IOInProgress - prof.prior.Device[i].IOInProgress
		u.Device[i].IOTime = cur.Device[i].IOTime - prof.prior.Device[i].IOTime
		u.Device[i].WeightedIOTime = cur.Device[i].WeightedIOTime - prof.prior.Device[i].WeightedIOTime
	}
	return u
}

// Ticker delivers the system's IO usage of the block devices at intervals.
type Ticker struct {
	*joe.Ticker
	Data chan *structs.DiskUsage
	*Profiler
}

// NewTicker returns a new Ticker containing a Data channel that delivers the
// data at intervals and an error channel that delivers any errors encountered.
// Stop the ticker to signal the ticker to stop running. Stopping the ticker
// does not close the Data channel; call Close to close both the ticker and the
// data channel.
func NewTicker(d time.Duration) (joe.Tocker, error) {
	p, err := NewProfiler()
	if err != nil {
		return nil, err
	}
	t := Ticker{Ticker: joe.NewTicker(d), Data: make(chan *structs.DiskUsage), Profiler: p}
	go t.Run()
	return &t, nil
}

// Run runs the ticker.
func (t *Ticker) Run() {
	var (
		i, priorPos, pos, line, fieldNum int
		n                                uint64
		v                                byte
		err                              error
		dev                              structs.Device
		cur                              structs.DiskStats
	)
	// ticker
	for {
		select {
		case <-t.Done:
			return
		case <-t.C:
			cur.Timestamp = time.Now().UTC().UnixNano()
			err = t.Reset()
			if err != nil {
				t.Errs <- err
				break
			}
			cur.Device = cur.Device[:0]
			// read each line until eof
			for {
				t.Val = t.Val[:0]
				t.Line, err = t.Buf.ReadSlice('\n')
				if err != nil {
					if err == io.EOF {
						break
					}
					t.Errs <- &joe.ReadError{Err: err}
					break
				}
				line++
				pos = 0
				fieldNum = 0
				// process the fields in the line
				for {
					// ignore spaces on the first two fields
					if fieldNum < 2 {
						for i, v = range t.Line[pos:] {
							if v != 0x20 {
								break
							}
						}
						pos += i
					}
					fieldNum++
					for i, v = range t.Line[pos:] {
						if v == 0x20 || v == '\n' {
							break
						}
					}
					if fieldNum != 3 {
						n, err = helpers.ParseUint(t.Line[pos : pos+i])
						if err != nil {
							t.Errs <- &joe.ParseError{Info: fmt.Sprintf("line %d: field %d", line, fieldNum), Err: err}
							continue
						}
					}
					priorPos, pos = pos, pos+i+1
					if fieldNum < 8 {
						if fieldNum < 4 {
							if fieldNum < 2 {
								if fieldNum == 1 {
									dev.Major = uint32(n)
									continue
								}
								dev.Minor = uint32(n)
								continue
							}
							dev.Name = string(t.Line[priorPos:pos])
							continue
						}
						if fieldNum < 6 {
							if fieldNum == 4 {
								dev.ReadsCompleted = n
								continue
							}
							dev.ReadsMerged = n
							continue
						}
						if fieldNum == 6 {
							dev.ReadSectors = n
							continue
						}
						dev.ReadingTime = n
						continue
					}
					if fieldNum < 12 {
						if fieldNum < 10 {
							if fieldNum == 8 {
								dev.WritesCompleted = n
								continue
							}
							dev.WritesMerged = n
							continue
						}
						if fieldNum == 10 {
							dev.WrittenSectors = n
							continue
						}
						dev.WritingTime = n
						continue
					}
					if fieldNum == 12 {
						dev.IOInProgress = int32(n)
						continue
					}
					if fieldNum == 13 {
						dev.IOTime = n
						continue
					}
					dev.WeightedIOTime = n
					break
				}
				cur.Device = append(cur.Device, dev)
			}
			t.Data <- t.CalculateUsage(&cur)
			// set prior info
			t.prior.Timestamp = cur.Timestamp
			if len(t.prior.Device) != len(cur.Device) {
				t.prior.Device = make([]structs.Device, len(cur.Device))
			}
			copy(t.prior.Device, cur.Device)
		}
	}
}

// Close closes the ticker resources.
func (t *Ticker) Close() {
	t.Ticker.Close()
	close(t.Data)
}
<commit_msg>Buf.ReadSlice() -> .ReadSlice() as it is now an interface method<commit_after>// Copyright 2016 Joel Scoble and The JoeFriday 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 diskusage calculates IO usage of the block devices. Usage is
// calculated by taking the difference between two snapshots of IO statistics
// for block devices, /procd/diskstats. The time elapsed between the two
// snapshots is stored in the TimeDelta field.
package diskusage

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

	"github.com/SermoDigital/helpers"
	joe "github.com/mohae/joefriday"
	stats "github.com/mohae/joefriday/disk/diskstats"
	"github.com/mohae/joefriday/disk/structs"
)

// Profiler is used to process the IO usage of the block devices.
type Profiler struct {
	*stats.Profiler
	prior *structs.DiskStats
}

// Returns an initialized Profiler; ready to use. Upon creation, a
// /proc/diskstats snapshot is taken so that any Get() will return valid
// information.
func NewProfiler() (prof *Profiler, err error) {
	p, err := stats.NewProfiler()
	if err != nil {
		return nil, err
	}
	s, err := p.Get()
	if err != nil {
		return nil, err
	}
	return &Profiler{Profiler: p, prior: s}, nil
}

// Get returns the current IO usage of the block devices. Calculating usage
// requires two snapshots. This func gets the current snapshot of
// /proc/diskstats and calculates the difference between that and the prior
// snapshot. The current snapshot is stored for use as the prior snapshot on
// the next Get call. If ongoing usage information is desired, the Ticker
// should be used; it's better suited for ongoing usage information.
func (prof *Profiler) Get() (u *structs.DiskUsage, err error) {
	st, err := prof.Profiler.Get()
	if err != nil {
		return nil, err
	}
	u = prof.CalculateUsage(st)
	prof.prior = st
	return u, nil
}

var std *Profiler
var stdMu sync.Mutex

// Get returns the current IO usage of the block devices using the package's
// global Profiler. The Profiler is instantiated lazily. If the profiler
// doesn't already exist, the first usage information will not be useful due
// to minimal time elapsing between the initial and second snapshots used for
// usage calculations; the results of the first call should be discarded.
func Get() (u *structs.DiskUsage, err error) {
	stdMu.Lock()
	defer stdMu.Unlock()
	if std == nil {
		std, err = NewProfiler()
		if err != nil {
			return nil, err
		}
	}
	return std.Get()
}

// CalculateUsage returns the difference between the current /proc/diskstats
// snapshot and the prior one.
func (prof *Profiler) CalculateUsage(cur *structs.DiskStats) *structs.DiskUsage {
	u := &structs.DiskUsage{Timestamp: cur.Timestamp, Device: make([]structs.Device, len(cur.Device))}
	u.TimeDelta = cur.Timestamp - prof.prior.Timestamp
	for i := 0; i < len(cur.Device); i++ {
		u.Device[i].Major = cur.Device[i].Major
		u.Device[i].Minor = cur.Device[i].Minor
		u.Device[i].Name = cur.Device[i].Name
		u.Device[i].ReadsCompleted = cur.Device[i].ReadsCompleted - prof.prior.Device[i].ReadsCompleted
		u.Device[i].ReadsMerged = cur.Device[i].ReadsMerged - prof.prior.Device[i].ReadsMerged
		u.Device[i].ReadSectors = cur.Device[i].ReadSectors - prof.prior.Device[i].ReadSectors
		u.Device[i].ReadingTime = cur.Device[i].ReadingTime - prof.prior.Device[i].ReadingTime
		u.Device[i].WritesCompleted = cur.Device[i].WritesCompleted - prof.prior.Device[i].WritesCompleted
		u.Device[i].WritesMerged = cur.Device[i].WritesMerged - prof.prior.Device[i].WritesMerged
		u.Device[i].WrittenSectors = cur.Device[i].WrittenSectors - prof.prior.Device[i].WrittenSectors
		u.Device[i].WritingTime = cur.Device[i].WritingTime - prof.prior.Device[i].WritingTime
		u.Device[i].IOInProgress = cur.Device[i].IOInProgress - prof.prior.Device[i].IOInProgress
		u.Device[i].IOTime = cur.Device[i].IOTime - prof.prior.Device[i].IOTime
		u.Device[i].WeightedIOTime = cur.Device[i].WeightedIOTime - prof.prior.Device[i].WeightedIOTime
	}
	return u
}

// Ticker delivers the system's IO usage of the block devices at intervals.
type Ticker struct {
	*joe.Ticker
	Data chan *structs.DiskUsage
	*Profiler
}

// NewTicker returns a new Ticker containing a Data channel that delivers the
// data at intervals and an error channel that delivers any errors encountered.
// Stop the ticker to signal the ticker to stop running. Stopping the ticker
// does not close the Data channel; call Close to close both the ticker and the
// data channel.
func NewTicker(d time.Duration) (joe.Tocker, error) {
	p, err := NewProfiler()
	if err != nil {
		return nil, err
	}
	t := Ticker{Ticker: joe.NewTicker(d), Data: make(chan *structs.DiskUsage), Profiler: p}
	go t.Run()
	return &t, nil
}

// Run runs the ticker.
func (t *Ticker) Run() {
	var (
		i, priorPos, pos, line, fieldNum int
		n                                uint64
		v                                byte
		err                              error
		dev                              structs.Device
		cur                              structs.DiskStats
	)
	// ticker
	for {
		select {
		case <-t.Done:
			return
		case <-t.C:
			cur.Timestamp = time.Now().UTC().UnixNano()
			err = t.Reset()
			if err != nil {
				t.Errs <- err
				break
			}
			cur.Device = cur.Device[:0]
			// read each line until eof
			for {
				t.Val = t.Val[:0]
				t.Line, err = t.ReadSlice('\n')
				if err != nil {
					if err == io.EOF {
						break
					}
					t.Errs <- &joe.ReadError{Err: err}
					break
				}
				line++
				pos = 0
				fieldNum = 0
				// process the fields in the line
				for {
					// ignore spaces on the first two fields
					if fieldNum < 2 {
						for i, v = range t.Line[pos:] {
							if v != 0x20 {
								break
							}
						}
						pos += i
					}
					fieldNum++
					for i, v = range t.Line[pos:] {
						if v == 0x20 || v == '\n' {
							break
						}
					}
					if fieldNum != 3 {
						n, err = helpers.ParseUint(t.Line[pos : pos+i])
						if err != nil {
							t.Errs <- &joe.ParseError{Info: fmt.Sprintf("line %d: field %d", line, fieldNum), Err: err}
							continue
						}
					}
					priorPos, pos = pos, pos+i+1
					if fieldNum < 8 {
						if fieldNum < 4 {
							if fieldNum < 2 {
								if fieldNum == 1 {
									dev.Major = uint32(n)
									continue
								}
								dev.Minor = uint32(n)
								continue
							}
							dev.Name = string(t.Line[priorPos:pos])
							continue
						}
						if fieldNum < 6 {
							if fieldNum == 4 {
								dev.ReadsCompleted = n
								continue
							}
							dev.ReadsMerged = n
							continue
						}
						if fieldNum == 6 {
							dev.ReadSectors = n
							continue
						}
						dev.ReadingTime = n
						continue
					}
					if fieldNum < 12 {
						if fieldNum < 10 {
							if fieldNum == 8 {
								dev.WritesCompleted = n
								continue
							}
							dev.WritesMerged = n
							continue
						}
						if fieldNum == 10 {
							dev.WrittenSectors = n
							continue
						}
						dev.WritingTime = n
						continue
					}
					if fieldNum == 12 {
						dev.IOInProgress = int32(n)
						continue
					}
					if fieldNum == 13 {
						dev.IOTime = n
						continue
					}
					dev.WeightedIOTime = n
					break
				}
				cur.Device = append(cur.Device, dev)
			}
			t.Data <- t.CalculateUsage(&cur)
			// set prior info
			t.prior.Timestamp = cur.Timestamp
			if len(t.prior.Device) != len(cur.Device) {
				t.prior.Device = make([]structs.Device, len(cur.Device))
			}
			copy(t.prior.Device, cur.Device)
		}
	}
}

// Close closes the ticker resources.
func (t *Ticker) Close() {
	t.Ticker.Close()
	close(t.Data)
}
<|endoftext|>


<commit_before>package torrentfs

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"

	"bitbucket.org/anacrolix/go.torrent"
	"bitbucket.org/anacrolix/go.torrent/testutil"
	"bitbucket.org/anacrolix/go.torrent/util"
	"github.com/anacrolix/libtorgo/metainfo"

	"bazil.org/fuse"
	fusefs "bazil.org/fuse/fs"
)

func init() {
	go http.ListenAndServe(":6061", nil)
}

func TestTCPAddrString(t *testing.T) {
	l, err := net.Listen("tcp4", "localhost:0")
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()
	c, err := net.Dial("tcp", l.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	ras := c.RemoteAddr().String()
	ta := &net.TCPAddr{
		IP:   net.IPv4(127, 0, 0, 1),
		Port: util.AddrPort(l.Addr()),
	}
	s := ta.String()
	if ras != s {
		t.FailNow()
	}
}

type testLayout struct {
	BaseDir   string
	MountDir  string
	Completed string
	Metainfo  *metainfo.MetaInfo
}

func (me *testLayout) Destroy() error {
	return os.RemoveAll(me.BaseDir)
}

func newGreetingLayout() (tl testLayout, err error) {
	tl.BaseDir, err = ioutil.TempDir("", "torrentfs")
	if err != nil {
		return
	}
	tl.Completed = filepath.Join(tl.BaseDir, "completed")
	os.Mkdir(tl.Completed, 0777)
	tl.MountDir = filepath.Join(tl.BaseDir, "mnt")
	os.Mkdir(tl.MountDir, 0777)
	name := testutil.CreateDummyTorrentData(tl.Completed)
	metaInfoBuf := &bytes.Buffer{}
	testutil.CreateMetaInfo(name, metaInfoBuf)
	tl.Metainfo, err = metainfo.Load(metaInfoBuf)
	log.Printf("%x", tl.Metainfo.Info.Pieces)
	return
}

// Unmount without first killing the FUSE connection while there are FUSE
// operations blocked inside the filesystem code.
func TestUnmountWedged(t *testing.T) {
	layout, err := newGreetingLayout()
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		err := layout.Destroy()
		if err != nil {
			t.Log(err)
		}
	}()
	client, err := torrent.NewClient(&torrent.Config{
		DataDir:         filepath.Join(layout.BaseDir, "incomplete"),
		DisableTrackers: true,
		NoDHT:           true,

		NoDefaultBlocklist: true,
	})
	defer client.Stop()
	t.Logf("%+v", *layout.Metainfo)
	client.AddTorrent(layout.Metainfo)
	fs := New(client)
	fuseConn, err := fuse.Mount(layout.MountDir)
	if err != nil {
		if strings.Contains(err.Error(), "fuse") {
			t.Skip(err)
		}
		t.Fatal(err)
	}
	<-fuseConn.Ready
	if err := fuseConn.MountError; err != nil {
		t.Fatal(err)
	}
	go func() {
		server := fusefs.Server{
			FS: fs,
			Debug: func(msg interface{}) {
				t.Log(msg)
			},
		}
		server.Serve(fuseConn)
	}()
	// Read the greeting file, though it will never be available. This should
	// "wedge" FUSE, requiring the fs object to be forcibly destroyed. The
	// read call will return with a FS error.
	go func() {
		_, err := ioutil.ReadFile(filepath.Join(layout.MountDir, layout.Metainfo.Info.Name))
		if err == nil {
			t.Fatal("expected error reading greeting")
		}
	}()

	// Wait until the read has blocked inside the filesystem code.
	fs.mu.Lock()
	for fs.blockedReads != 1 {
		fs.event.Wait()
	}
	fs.mu.Unlock()

	fs.Destroy()

	for {
		err = fuse.Unmount(layout.MountDir)
		if err != nil {
			t.Logf("error unmounting: %s", err)
		} else {
			break
		}
	}

	err = fuseConn.Close()
	if err != nil {
		t.Fatalf("error closing fuse conn: %s", err)
	}
}

func TestDownloadOnDemand(t *testing.T) {
	layout, err := newGreetingLayout()
	if err != nil {
		t.Fatal(err)
	}
	defer layout.Destroy()
	seeder, err := torrent.NewClient(&torrent.Config{
		DataDir:         layout.Completed,
		DisableTrackers: true,
		NoDHT:           true,
		ListenAddr:      ":0",

		NoDefaultBlocklist: true,
	})
	if err != nil {
		t.Fatalf("error creating seeder client: %s", err)
	}
	seeder.SetIPBlockList(nil)
	defer seeder.Stop()
	http.HandleFunc("/seeder", func(w http.ResponseWriter, req *http.Request) {
		seeder.WriteStatus(w)
	})
	_, err = seeder.AddMagnet(fmt.Sprintf("magnet:?xt=urn:btih:%x", layout.Metainfo.Info.Hash))
	if err != nil {
		t.Fatal(err)
	}
	leecher, err := torrent.NewClient(&torrent.Config{
		DataDir:         filepath.Join(layout.BaseDir, "download"),
		DisableTrackers: true,
		NoDHT:           true,
		ListenAddr:      ":0",

		NoDefaultBlocklist: true,

		// This can be used to check if clients can connect to other clients
		// with the same ID.

		// PeerID: seeder.PeerID(),
	})
	leecher.SetIPBlockList(nil)
	http.HandleFunc("/leecher", func(w http.ResponseWriter, req *http.Request) {
		leecher.WriteStatus(w)
	})
	defer leecher.Stop()
	leecher.AddTorrent(layout.Metainfo)
	var ih torrent.InfoHash
	util.CopyExact(ih[:], layout.Metainfo.Info.Hash)
	leecher.AddPeers(ih, []torrent.Peer{func() torrent.Peer {
		_, port, err := net.SplitHostPort(seeder.ListenAddr().String())
		if err != nil {
			panic(err)
		}
		portInt64, err := strconv.ParseInt(port, 0, 0)
		if err != nil {
			panic(err)
		}
		return torrent.Peer{
			IP:   net.IPv6loopback,
			Port: int(portInt64),
		}
	}()})
	fs := New(leecher)
	defer fs.Destroy()
	root, _ := fs.Root()
	node, _ := root.(fusefs.NodeStringLookuper).Lookup("greeting", nil)
	size := int(node.Attr().Size)
	resp := &fuse.ReadResponse{
		Data: make([]byte, size),
	}
	node.(fusefs.HandleReader).Read(&fuse.ReadRequest{
		Size: size,
	}, resp, nil)
	content := resp.Data
	if string(content) != testutil.GreetingFileContents {
		t.FailNow()
	}
}
<commit_msg>Fix TestUnmountWedged test on OSX<commit_after>package torrentfs

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	_ "net/http/pprof"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"
	"time"

	"bitbucket.org/anacrolix/go.torrent"
	"bitbucket.org/anacrolix/go.torrent/testutil"
	"bitbucket.org/anacrolix/go.torrent/util"
	"github.com/anacrolix/libtorgo/metainfo"

	"bazil.org/fuse"
	fusefs "bazil.org/fuse/fs"
)

func init() {
	go http.ListenAndServe(":6061", nil)
}

func TestTCPAddrString(t *testing.T) {
	l, err := net.Listen("tcp4", "localhost:0")
	if err != nil {
		t.Fatal(err)
	}
	defer l.Close()
	c, err := net.Dial("tcp", l.Addr().String())
	if err != nil {
		t.Fatal(err)
	}
	defer c.Close()
	ras := c.RemoteAddr().String()
	ta := &net.TCPAddr{
		IP:   net.IPv4(127, 0, 0, 1),
		Port: util.AddrPort(l.Addr()),
	}
	s := ta.String()
	if ras != s {
		t.FailNow()
	}
}

type testLayout struct {
	BaseDir   string
	MountDir  string
	Completed string
	Metainfo  *metainfo.MetaInfo
}

func (me *testLayout) Destroy() error {
	return os.RemoveAll(me.BaseDir)
}

func newGreetingLayout() (tl testLayout, err error) {
	tl.BaseDir, err = ioutil.TempDir("", "torrentfs")
	if err != nil {
		return
	}
	tl.Completed = filepath.Join(tl.BaseDir, "completed")
	os.Mkdir(tl.Completed, 0777)
	tl.MountDir = filepath.Join(tl.BaseDir, "mnt")
	os.Mkdir(tl.MountDir, 0777)
	name := testutil.CreateDummyTorrentData(tl.Completed)
	metaInfoBuf := &bytes.Buffer{}
	testutil.CreateMetaInfo(name, metaInfoBuf)
	tl.Metainfo, err = metainfo.Load(metaInfoBuf)
	return
}

// Unmount without first killing the FUSE connection while there are FUSE
// operations blocked inside the filesystem code.
func TestUnmountWedged(t *testing.T) {
	layout, err := newGreetingLayout()
	if err != nil {
		t.Fatal(err)
	}
	defer func() {
		err := layout.Destroy()
		if err != nil {
			t.Log(err)
		}
	}()
	client, err := torrent.NewClient(&torrent.Config{
		DataDir:         filepath.Join(layout.BaseDir, "incomplete"),
		DisableTrackers: true,
		NoDHT:           true,

		NoDefaultBlocklist: true,
	})
	defer client.Stop()
	client.AddTorrent(layout.Metainfo)
	fs := New(client)
	fuseConn, err := fuse.Mount(layout.MountDir)
	if err != nil {
		if strings.Contains(err.Error(), "fuse") {
			t.Skip(err)
		}
		t.Fatal(err)
	}
	go func() {
		server := fusefs.Server{
			FS: fs,
			Debug: func(msg interface{}) {
				t.Log(msg)
			},
		}
		server.Serve(fuseConn)
	}()
	<-fuseConn.Ready
	if err := fuseConn.MountError; err != nil {
		t.Fatal(err)
	}
	// Read the greeting file, though it will never be available. This should
	// "wedge" FUSE, requiring the fs object to be forcibly destroyed. The
	// read call will return with a FS error.
	go func() {
		_, err := ioutil.ReadFile(filepath.Join(layout.MountDir, layout.Metainfo.Info.Name))
		if err == nil {
			t.Fatal("expected error reading greeting")
		}
	}()

	// Wait until the read has blocked inside the filesystem code.
	fs.mu.Lock()
	for fs.blockedReads != 1 {
		fs.event.Wait()
	}
	fs.mu.Unlock()

	fs.Destroy()

	for {
		err = fuse.Unmount(layout.MountDir)
		if err != nil {
			t.Logf("error unmounting: %s", err)
			time.Sleep(time.Millisecond)
		} else {
			break
		}
	}

	err = fuseConn.Close()
	if err != nil {
		t.Fatalf("error closing fuse conn: %s", err)
	}
}

func TestDownloadOnDemand(t *testing.T) {
	layout, err := newGreetingLayout()
	if err != nil {
		t.Fatal(err)
	}
	defer layout.Destroy()
	seeder, err := torrent.NewClient(&torrent.Config{
		DataDir:         layout.Completed,
		DisableTrackers: true,
		NoDHT:           true,
		ListenAddr:      ":0",

		NoDefaultBlocklist: true,
	})
	if err != nil {
		t.Fatalf("error creating seeder client: %s", err)
	}
	seeder.SetIPBlockList(nil)
	defer seeder.Stop()
	http.HandleFunc("/seeder", func(w http.ResponseWriter, req *http.Request) {
		seeder.WriteStatus(w)
	})
	_, err = seeder.AddMagnet(fmt.Sprintf("magnet:?xt=urn:btih:%x", layout.Metainfo.Info.Hash))
	if err != nil {
		t.Fatal(err)
	}
	leecher, err := torrent.NewClient(&torrent.Config{
		DataDir:         filepath.Join(layout.BaseDir, "download"),
		DisableTrackers: true,
		NoDHT:           true,
		ListenAddr:      ":0",

		NoDefaultBlocklist: true,

		// This can be used to check if clients can connect to other clients
		// with the same ID.

		// PeerID: seeder.PeerID(),
	})
	leecher.SetIPBlockList(nil)
	http.HandleFunc("/leecher", func(w http.ResponseWriter, req *http.Request) {
		leecher.WriteStatus(w)
	})
	defer leecher.Stop()
	leecher.AddTorrent(layout.Metainfo)
	var ih torrent.InfoHash
	util.CopyExact(ih[:], layout.Metainfo.Info.Hash)
	leecher.AddPeers(ih, []torrent.Peer{func() torrent.Peer {
		_, port, err := net.SplitHostPort(seeder.ListenAddr().String())
		if err != nil {
			panic(err)
		}
		portInt64, err := strconv.ParseInt(port, 0, 0)
		if err != nil {
			panic(err)
		}
		return torrent.Peer{
			IP:   net.IPv6loopback,
			Port: int(portInt64),
		}
	}()})
	fs := New(leecher)
	defer fs.Destroy()
	root, _ := fs.Root()
	node, _ := root.(fusefs.NodeStringLookuper).Lookup("greeting", nil)
	size := int(node.Attr().Size)
	resp := &fuse.ReadResponse{
		Data: make([]byte, size),
	}
	node.(fusefs.HandleReader).Read(&fuse.ReadRequest{
		Size: size,
	}, resp, nil)
	content := resp.Data
	if string(content) != testutil.GreetingFileContents {
		t.FailNow()
	}
}
<|endoftext|>


<commit_before>// Package function implements function-level opterations.
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 generates a zip and creates or updates the function.
func (f *Function) Deploy() 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>fix typo in comment<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 generates a zip and creates or updates the function.
func (f *Function) Deploy() 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
}
<|endoftext|>