<commit_before>package main

import (
	"fmt"
	"log"
	"math/rand"
	"os"
	"time"

	"gopkg.in/mgo.v2/bson"
)

var (
	botStartTime time.Time
	botResponses int
	botMessages  int
)

func init() {
	rand.Seed(time.Now().Unix())
	botStartTime = time.Now()
}

func getUptime() string {
	return time.Now().Sub(botStartTime).String()
}

//getInfo returns the bot infromation as a discord formated string
func getInfo() (result string) {
	result += "```"
	result += fmt.Sprintf("Name: %s\n", discordCfg.Name)
	result += fmt.Sprintf("Uptime: %s\n", getUptime())
	result += fmt.Sprintf("Guilds: %d\n", len(discordCfg.Guilds))
	result += fmt.Sprintf("Anime channels: %d\n", len(discordCfg.AnimeChannels))
	result += fmt.Sprintf("Unique subscribers: %d\n", getUniqueSubs())
	result += fmt.Sprintf("Messages read: %d\n", botMessages)
	result += fmt.Sprintf("Message responses: %d\n", botResponses)
	result += "```"

	return
}

//getUniqueSubs returs the amount of unique subscribers in the database
func getUniqueSubs() int {
	var uniqueSubs []string
	err := DBanimeList.Find(bson.M{}).Distinct("subs", &uniqueSubs)
	if err != nil {
		log.Print("MongoDB error: ", err)
	}

	return len(uniqueSubs)
}

//appendUnique is a function for appending a string to string array
//only appends the string if it doesn't already exsist in the array
func appendUnique(slice []string, id string) []string {
	for _, s := range slice {
		if s == id {
			return slice
		}
	}
	return append(slice, id)
}

//removeItem removes a string from the array and returns the new array
//first param is the array to remove from
//second param is the string to remove from array (removes all instances)
func removeItem(slice []string, item string) []string {
	for i, value := range slice {
		if value == item {
			slice = append(slice[:i], slice[i+1:]...)
		}
	}
	return slice
}

//checks if any of the provided required ENV variables is missing
//logs the error and exits the program if any of them is missing
func requireEnvVars(args ...string) {
	for _, v := range args {
		if os.Getenv(v) == "" {
			log.Fatal("Missing ENV variable: ", v)
		}
	}
}
<commit_msg>fixed comment<commit_after>package main

import (
	"fmt"
	"log"
	"math/rand"
	"os"
	"time"

	"gopkg.in/mgo.v2/bson"
)

var (
	botStartTime time.Time
	botResponses int
	botMessages  int
)

func init() {
	rand.Seed(time.Now().Unix())
	botStartTime = time.Now()
}

func getUptime() string {
	return time.Now().Sub(botStartTime).String()
}

//getInfo returns the bot infromation as a discord formated string
func getInfo() (result string) {
	result += "```"
	result += fmt.Sprintf("Name: %s\n", discordCfg.Name)
	result += fmt.Sprintf("Uptime: %s\n", getUptime())
	result += fmt.Sprintf("Guilds: %d\n", len(discordCfg.Guilds))
	result += fmt.Sprintf("Anime channels: %d\n", len(discordCfg.AnimeChannels))
	result += fmt.Sprintf("Unique subscribers: %d\n", getUniqueSubs())
	result += fmt.Sprintf("Messages read: %d\n", botMessages)
	result += fmt.Sprintf("Message responses: %d\n", botResponses)
	result += "```"

	return
}

//getUniqueSubs returs the amount of unique subscribers in the database
func getUniqueSubs() int {
	var uniqueSubs []string
	err := DBanimeList.Find(bson.M{}).Distinct("subs", &uniqueSubs)
	if err != nil {
		log.Print("MongoDB error: ", err)
	}

	return len(uniqueSubs)
}

//appendUnique is a function for appending a string to string array
//only appends the string if it doesn't already exsist in the array
func appendUnique(slice []string, id string) []string {
	for _, s := range slice {
		if s == id {
			return slice
		}
	}
	return append(slice, id)
}

//removeItem removes a string from the array and returns the new array
//first param is the array to remove from
//second param is the string to remove from array (removes all instances)
func removeItem(slice []string, item string) []string {
	for i, value := range slice {
		if value == item {
			slice = append(slice[:i], slice[i+1:]...)
		}
	}
	return slice
}

// requireEnvVars checks if any of the required ENV variables are missing
// logs the error and exits the program if any of them are missing
func requireEnvVars(args ...string) {
	for _, v := range args {
		if os.Getenv(v) == "" {
			log.Fatal("Missing ENV variable: ", v)
		}
	}
}
<|endoftext|>


<commit_before>package sarama

import (
	"fmt"
	"testing"
)

func ExampleBroker() error {
	broker := NewBroker("localhost:9092")
	err := broker.Open(nil)
	if err != nil {
		return err
	}

	request := MetadataRequest{Topics: []string{"myTopic"}}
	response, err := broker.GetMetadata(&request)
	if err != nil {
		_ = broker.Close()
		return err
	}

	fmt.Println("There are", len(response.Topics), "topics active in the cluster.")

	return broker.Close()
}

type mockEncoder struct {
	bytes []byte
}

func (m mockEncoder) encode(pe packetEncoder) error {
	return pe.putRawBytes(m.bytes)
}

func TestBrokerAccessors(t *testing.T) {
	broker := NewBroker("abc:123")

	if broker.ID() != -1 {
		t.Error("New broker didn't have an ID of -1.")
	}

	if broker.Addr() != "abc:123" {
		t.Error("New broker didn't have the correct address")
	}

	broker.id = 34
	if broker.ID() != 34 {
		t.Error("Manually setting broker ID did not take effect.")
	}
}

func TestSimpleBrokerCommunication(t *testing.T) {
	mb := newMockBroker(t, 0)
	defer mb.Close()

	broker := NewBroker(mb.Addr())
	err := broker.Open(nil)
	if err != nil {
		t.Fatal(err)
	}

	for _, tt := range brokerTestTable {
		mb.Returns(&mockEncoder{tt.response})
	}
	for _, tt := range brokerTestTable {
		tt.runner(t, broker)
	}

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

// We're not testing encoding/decoding here, so most of the requests/responses will be empty for simplicity's sake
var brokerTestTable = []struct {
	response []byte
	runner   func(*testing.T, *Broker)
}{
	{[]byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := MetadataRequest{}
			response, err := broker.GetMetadata(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("Metadata request got no response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 't', 0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := ConsumerMetadataRequest{}
			response, err := broker.GetConsumerMetadata(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("Consumer Metadata request got no response!")
			}
		}},

	{[]byte{},
		func(t *testing.T, broker *Broker) {
			request := ProduceRequest{}
			request.RequiredAcks = NoResponse
			response, err := broker.Produce(&request)
			if err != nil {
				t.Error(err)
			}
			if response != nil {
				t.Error("Produce request with NoResponse got a response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := ProduceRequest{}
			request.RequiredAcks = WaitForLocal
			response, err := broker.Produce(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("Produce request without NoResponse got no response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := FetchRequest{}
			response, err := broker.Fetch(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("Fetch request got no response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := OffsetFetchRequest{}
			response, err := broker.FetchOffset(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("OffsetFetch request got no response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := OffsetCommitRequest{}
			response, err := broker.CommitOffset(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("OffsetCommit request got no response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := OffsetRequest{}
			response, err := broker.GetAvailableOffsets(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("Offset request got no response!")
			}
		}},
}
<commit_msg>Don't return from example functions<commit_after>package sarama

import (
	"fmt"
	"testing"
)

func ExampleBroker() {
	broker := NewBroker("localhost:9092")
	err := broker.Open(nil)
	if err != nil {
		panic(err)
	}

	request := MetadataRequest{Topics: []string{"myTopic"}}
	response, err := broker.GetMetadata(&request)
	if err != nil {
		_ = broker.Close()
		panic(err)
	}

	fmt.Println("There are", len(response.Topics), "topics active in the cluster.")

	if err = broker.Close(); err != nil {
		panic(err)
	}
}

type mockEncoder struct {
	bytes []byte
}

func (m mockEncoder) encode(pe packetEncoder) error {
	return pe.putRawBytes(m.bytes)
}

func TestBrokerAccessors(t *testing.T) {
	broker := NewBroker("abc:123")

	if broker.ID() != -1 {
		t.Error("New broker didn't have an ID of -1.")
	}

	if broker.Addr() != "abc:123" {
		t.Error("New broker didn't have the correct address")
	}

	broker.id = 34
	if broker.ID() != 34 {
		t.Error("Manually setting broker ID did not take effect.")
	}
}

func TestSimpleBrokerCommunication(t *testing.T) {
	mb := newMockBroker(t, 0)
	defer mb.Close()

	broker := NewBroker(mb.Addr())
	err := broker.Open(nil)
	if err != nil {
		t.Fatal(err)
	}

	for _, tt := range brokerTestTable {
		mb.Returns(&mockEncoder{tt.response})
	}
	for _, tt := range brokerTestTable {
		tt.runner(t, broker)
	}

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

// We're not testing encoding/decoding here, so most of the requests/responses will be empty for simplicity's sake
var brokerTestTable = []struct {
	response []byte
	runner   func(*testing.T, *Broker)
}{
	{[]byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := MetadataRequest{}
			response, err := broker.GetMetadata(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("Metadata request got no response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 't', 0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := ConsumerMetadataRequest{}
			response, err := broker.GetConsumerMetadata(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("Consumer Metadata request got no response!")
			}
		}},

	{[]byte{},
		func(t *testing.T, broker *Broker) {
			request := ProduceRequest{}
			request.RequiredAcks = NoResponse
			response, err := broker.Produce(&request)
			if err != nil {
				t.Error(err)
			}
			if response != nil {
				t.Error("Produce request with NoResponse got a response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := ProduceRequest{}
			request.RequiredAcks = WaitForLocal
			response, err := broker.Produce(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("Produce request without NoResponse got no response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := FetchRequest{}
			response, err := broker.Fetch(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("Fetch request got no response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := OffsetFetchRequest{}
			response, err := broker.FetchOffset(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("OffsetFetch request got no response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := OffsetCommitRequest{}
			response, err := broker.CommitOffset(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("OffsetCommit request got no response!")
			}
		}},

	{[]byte{0x00, 0x00, 0x00, 0x00},
		func(t *testing.T, broker *Broker) {
			request := OffsetRequest{}
			response, err := broker.GetAvailableOffsets(&request)
			if err != nil {
				t.Error(err)
			}
			if response == nil {
				t.Error("Offset request got no response!")
			}
		}},
}
<|endoftext|>


<commit_before>package worker

import (
	"encoding/json"
	"fmt"
	"github.com/go-redis/redis"
	"github.com/stitchfix/flotilla-os/clients/metrics"
	"github.com/stitchfix/flotilla-os/queue"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/execution/engine"
	flotillaLog "github.com/stitchfix/flotilla-os/log"
	"github.com/stitchfix/flotilla-os/state"
	"gopkg.in/tomb.v2"
)

type statusWorker struct {
	sm                       state.Manager
	ee                       engine.Engine
	conf                     config.Config
	log                      flotillaLog.Logger
	pollInterval             time.Duration
	t                        tomb.Tomb
	engine                   *string
	redisClient              *redis.Client
	workerId                 string
	exceptionExtractorClient *http.Client
	exceptionExtractorUrl    string
}

func (sw *statusWorker) Initialize(conf config.Config, sm state.Manager, ee engine.Engine, log flotillaLog.Logger, pollInterval time.Duration, engine *string, qm queue.Manager) error {
	sw.pollInterval = pollInterval
	sw.conf = conf
	sw.sm = sm
	sw.ee = ee
	sw.log = log
	sw.engine = engine
	sw.workerId = fmt.Sprintf("workerid:%d", rand.Int())

	if sw.conf.IsSet("eks.exception_extractor_url") {
		sw.exceptionExtractorClient = &http.Client{
			Timeout: time.Second * 5,
		}
		sw.exceptionExtractorUrl = sw.conf.GetString("eks.exception_extractor_url")
	}
	sw.setupRedisClient(conf)
	_ = sw.log.Log("message", "initialized a status worker", "engine", *engine)
	return nil
}

func (sw *statusWorker) setupRedisClient(conf config.Config) {
	if *sw.engine == state.EKSEngine {
		sw.redisClient = redis.NewClient(&redis.Options{Addr: conf.GetString("redis_address"), DB: conf.GetInt("redis_db")})
	}
}

func (sw *statusWorker) GetTomb() *tomb.Tomb {
	return &sw.t
}

//
// Run updates status of tasks
//
func (sw *statusWorker) Run() error {
	for {
		select {
		case <-sw.t.Dying():
			sw.log.Log("message", "A status worker was terminated")
			return nil
		default:
			if *sw.engine == state.ECSEngine {
				sw.runOnceECS()
				time.Sleep(sw.pollInterval)
			}

			if *sw.engine == state.EKSEngine {
				sw.runOnceEKS()
				time.Sleep(sw.pollInterval)
			}
		}
	}
}

func (sw *statusWorker) runOnceEKS() {
	rl, err := sw.sm.ListRuns(1000, 0, "started_at", "asc", map[string][]string{
		"queued_at_since": {
			time.Now().AddDate(0, 0, -3).Format(time.RFC3339),
		},
		"status": {state.StatusNeedsRetry, state.StatusRunning, state.StatusQueued, state.StatusPending},
	}, nil, []string{state.EKSEngine})

	if err != nil {
		_ = sw.log.Log("message", "unable to receive runs", "error", fmt.Sprintf("%+v", err))
		return
	}
	runs := rl.Runs
	sw.processEKSRuns(runs)
}

func (sw *statusWorker) processEKSRuns(runs []state.Run) {
	var lockedRuns []state.Run
	for _, run := range runs {
		duration := time.Duration((rand.Intn(15))+25) * time.Second
		lock := sw.acquireLock(run, "status", duration)
		if lock {
			lockedRuns = append(lockedRuns, run)
		}
	}
	_ = metrics.Increment(metrics.StatusWorkerLockedRuns, []string{sw.workerId}, float64(len(lockedRuns)))
	for _, run := range lockedRuns {
		start := time.Now()
		sw.processEKSRun(run)
		_ = metrics.Timing(metrics.StatusWorkerProcessEKSRun, time.Since(start), []string{sw.workerId}, 1)
	}
}
func (sw *statusWorker) acquireLock(run state.Run, purpose string, expiration time.Duration) bool {
	start := time.Now()
	set, err := sw.redisClient.SetNX(fmt.Sprintf("%s-%s", run.RunID, purpose), sw.workerId, expiration).Result()
	if err != nil {
		_ = sw.log.Log("message", "unable to set lock", "error", fmt.Sprintf("%+v", err))
		return true
	}
	_ = metrics.Timing(metrics.StatusWorkerAcquireLock, time.Since(start), []string{sw.workerId}, 1)
	return set
}

func (sw *statusWorker) processEKSRun(run state.Run) {
	reloadRun, err := sw.sm.GetRun(run.RunID)
	if err == nil && reloadRun.Status == state.StatusStopped {
		// Run was updated by another worker process.
		return
	}
	start := time.Now()
	updatedRunWithMetrics, _ := sw.ee.FetchPodMetrics(run)
	_ = metrics.Timing(metrics.StatusWorkerFetchPodMetrics, time.Since(start), []string{sw.workerId}, 1)

	start = time.Now()
	updatedRun, err := sw.ee.FetchUpdateStatus(updatedRunWithMetrics)
	_ = metrics.Timing(metrics.StatusWorkerFetchUpdateStatus, time.Since(start), []string{sw.workerId}, 1)

	if err == nil {
		subRuns, err := sw.sm.ListRuns(1000, 0, "status", "desc", nil, map[string]string{"PARENT_FLOTILLA_RUN_ID": run.RunID}, state.Engines)
		if err == nil && subRuns.Total > 0 {
			var spawnedRuns state.SpawnedRuns
			for _, subRun := range subRuns.Runs {
				spawnedRuns = append(spawnedRuns, state.SpawnedRun{RunID: subRun.RunID})
			}
			updatedRun.SpawnedRuns = &spawnedRuns
		}
	}
	if err != nil {
		message := fmt.Sprintf("%+v", err)
		_ = sw.log.Log("message", "unable to receive eks runs", "error", message)

		minutesInQueue := time.Now().Sub(*run.QueuedAt).Minutes()
		if strings.Contains(message, "not found") && minutesInQueue > float64(30) {
			stoppedAt := time.Now()
			reason := "Job either timed out or not found on the EKS cluster."
			updatedRun.Status = state.StatusStopped
			updatedRun.FinishedAt = &stoppedAt
			updatedRun.ExitReason = &reason
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}

	} else {
		if run.Status != updatedRun.Status && (updatedRun.PodName == run.PodName) {
			_ = sw.log.Log("message", "updating eks run status", "pod", updatedRun.PodName, "status", updatedRun.Status, "exit_code", updatedRun.ExitCode)

			if sw.exceptionExtractorClient != nil && updatedRun.ExitCode != nil && *(updatedRun.ExitCode) != 0 {
				jobUrl := fmt.Sprintf("%s/extract/%s", sw.exceptionExtractorUrl, run.RunID)
				res, err := sw.exceptionExtractorClient.Get(jobUrl)
				if err == nil {
					body, err := ioutil.ReadAll(res.Body)
					if body != nil {
						defer res.Body.Close()
						runExceptions := state.RunExceptions{}
						err = json.Unmarshal(body, &runExceptions)
						if err == nil {
							updatedRun.RunExceptions = &runExceptions
						}
					}
				}
			}
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			if err != nil {
				_ = sw.log.Log("message", "unable to save eks runs", "error", fmt.Sprintf("%+v", err))
			}

			if updatedRun.Status == state.StatusStopped {
				//TODO - move to a separate worker.
				//_ = sw.ee.Terminate(run)
			}
		} else {
			if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
				updatedRun.MaxCpuUsed != run.MaxCpuUsed ||
				updatedRun.Cpu != run.Cpu ||
				updatedRun.Memory != run.Memory ||
				updatedRun.PodEvents != run.PodEvents ||
				updatedRun.SpawnedRuns != run.SpawnedRuns {
				_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			}
		}
	}
}

func (sw *statusWorker) processEKSRunMetrics(run state.Run) {
	updatedRun, err := sw.ee.FetchPodMetrics(run)
	if err == nil {
		if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
			updatedRun.MaxCpuUsed != run.MaxCpuUsed {
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}
	}
}

func (sw *statusWorker) runOnceECS() {
	runReceipt, err := sw.ee.PollStatus()

	if err != nil {
		sw.log.Log("message", "unable to receive status message", "error", fmt.Sprintf("%+v", err))
		return
	}

	// Ensure update is in the env required, otherwise, ack without taking action
	update := runReceipt.Run
	if update != nil {
		//
		// Relies on the reserved env var, FLOTILLA_SERVER_MODE to ensure update
		// belongs to -this- mode of Flotilla
		//
		var serverMode string
		if update.Env != nil {
			for _, kv := range *update.Env {
				if kv.Name == "FLOTILLA_SERVER_MODE" {
					serverMode = kv.Value
				}
			}
		}

		shouldProcess := len(serverMode) > 0 && serverMode == sw.conf.GetString("flotilla_mode")
		if shouldProcess {
			run, err := sw.findRun(update.TaskArn)
			if err != nil {
				sw.log.Log("message", "unable to find run to apply update to", "error", fmt.Sprintf("%+v", err))
				return
			}

			_, err = sw.sm.UpdateRun(run.RunID, *update)
			if err != nil {
				sw.log.Log("message", "error applying status update", "run", run.RunID, "error", fmt.Sprintf("%+v", err))
				return
			}

			// emit status update event
			sw.logStatusUpdate(*update)
		}

		if sw.conf.GetString("flotilla_mode") != "dev" {
			_ = sw.log.Log("message", "Acking status update", "arn", update.TaskArn)
		}
		if err = runReceipt.Done(); err != nil {
			sw.log.Log("message", "Acking status update failed", "arn", update.TaskArn, "error", fmt.Sprintf("%+v", err))
		}
	}
}

func (sw *statusWorker) logStatusUpdate(update state.Run) {
	var err error
	var startedAt, finishedAt time.Time
	var env state.EnvList

	if update.StartedAt != nil {
		startedAt = *update.StartedAt
	}

	if update.FinishedAt != nil {
		finishedAt = *update.FinishedAt
	}

	if update.Env != nil {
		env = *update.Env
	}

	if update.ExitCode != nil {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"exit_code", *update.ExitCode,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	} else {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	}

	if err != nil {
		sw.log.Log("message", "Failed to emit status event", "run_id", update.RunID, "error", err.Error())
	}
}

func (sw *statusWorker) findRun(taskArn string) (state.Run, error) {
	var engines []string
	if sw.engine != nil {
		engines = []string{*sw.engine}
	} else {
		engines = nil
	}

	runs, err := sw.sm.ListRuns(1, 0, "started_at", "asc", map[string][]string{
		"task_arn": {taskArn},
	}, nil, engines)
	if err != nil {
		return state.Run{}, errors.Wrapf(err, "problem finding run by task arn [%s]", taskArn)
	}
	if runs.Total > 0 && len(runs.Runs) > 0 {
		return runs.Runs[0], nil
	}
	return state.Run{}, errors.Errorf("no run found for [%s]", taskArn)
}
<commit_msg>only extract exceptions when exit code is not 0<commit_after>package worker

import (
	"encoding/json"
	"fmt"
	"github.com/go-redis/redis"
	"github.com/stitchfix/flotilla-os/clients/metrics"
	"github.com/stitchfix/flotilla-os/queue"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/stitchfix/flotilla-os/config"
	"github.com/stitchfix/flotilla-os/execution/engine"
	flotillaLog "github.com/stitchfix/flotilla-os/log"
	"github.com/stitchfix/flotilla-os/state"
	"gopkg.in/tomb.v2"
)

type statusWorker struct {
	sm                       state.Manager
	ee                       engine.Engine
	conf                     config.Config
	log                      flotillaLog.Logger
	pollInterval             time.Duration
	t                        tomb.Tomb
	engine                   *string
	redisClient              *redis.Client
	workerId                 string
	exceptionExtractorClient *http.Client
	exceptionExtractorUrl    string
}

func (sw *statusWorker) Initialize(conf config.Config, sm state.Manager, ee engine.Engine, log flotillaLog.Logger, pollInterval time.Duration, engine *string, qm queue.Manager) error {
	sw.pollInterval = pollInterval
	sw.conf = conf
	sw.sm = sm
	sw.ee = ee
	sw.log = log
	sw.engine = engine
	sw.workerId = fmt.Sprintf("workerid:%d", rand.Int())

	if sw.conf.IsSet("eks.exception_extractor_url") {
		sw.exceptionExtractorClient = &http.Client{
			Timeout: time.Second * 5,
		}
		sw.exceptionExtractorUrl = sw.conf.GetString("eks.exception_extractor_url")
	}
	sw.setupRedisClient(conf)
	_ = sw.log.Log("message", "initialized a status worker", "engine", *engine)
	return nil
}

func (sw *statusWorker) setupRedisClient(conf config.Config) {
	if *sw.engine == state.EKSEngine {
		sw.redisClient = redis.NewClient(&redis.Options{Addr: conf.GetString("redis_address"), DB: conf.GetInt("redis_db")})
	}
}

func (sw *statusWorker) GetTomb() *tomb.Tomb {
	return &sw.t
}

//
// Run updates status of tasks
//
func (sw *statusWorker) Run() error {
	for {
		select {
		case <-sw.t.Dying():
			sw.log.Log("message", "A status worker was terminated")
			return nil
		default:
			if *sw.engine == state.ECSEngine {
				sw.runOnceECS()
				time.Sleep(sw.pollInterval)
			}

			if *sw.engine == state.EKSEngine {
				sw.runOnceEKS()
				time.Sleep(sw.pollInterval)
			}
		}
	}
}

func (sw *statusWorker) runOnceEKS() {
	rl, err := sw.sm.ListRuns(1000, 0, "started_at", "asc", map[string][]string{
		"queued_at_since": {
			time.Now().AddDate(0, 0, -3).Format(time.RFC3339),
		},
		"status": {state.StatusNeedsRetry, state.StatusRunning, state.StatusQueued, state.StatusPending},
	}, nil, []string{state.EKSEngine})

	if err != nil {
		_ = sw.log.Log("message", "unable to receive runs", "error", fmt.Sprintf("%+v", err))
		return
	}
	runs := rl.Runs
	sw.processEKSRuns(runs)
}

func (sw *statusWorker) processEKSRuns(runs []state.Run) {
	var lockedRuns []state.Run
	for _, run := range runs {
		duration := time.Duration((rand.Intn(15))+25) * time.Second
		lock := sw.acquireLock(run, "status", duration)
		if lock {
			lockedRuns = append(lockedRuns, run)
		}
	}
	_ = metrics.Increment(metrics.StatusWorkerLockedRuns, []string{sw.workerId}, float64(len(lockedRuns)))
	for _, run := range lockedRuns {
		start := time.Now()
		sw.processEKSRun(run)
		_ = metrics.Timing(metrics.StatusWorkerProcessEKSRun, time.Since(start), []string{sw.workerId}, 1)
	}
}
func (sw *statusWorker) acquireLock(run state.Run, purpose string, expiration time.Duration) bool {
	start := time.Now()
	set, err := sw.redisClient.SetNX(fmt.Sprintf("%s-%s", run.RunID, purpose), sw.workerId, expiration).Result()
	if err != nil {
		_ = sw.log.Log("message", "unable to set lock", "error", fmt.Sprintf("%+v", err))
		return true
	}
	_ = metrics.Timing(metrics.StatusWorkerAcquireLock, time.Since(start), []string{sw.workerId}, 1)
	return set
}

func (sw *statusWorker) processEKSRun(run state.Run) {
	reloadRun, err := sw.sm.GetRun(run.RunID)
	if err == nil && reloadRun.Status == state.StatusStopped {
		// Run was updated by another worker process.
		return
	}
	start := time.Now()
	updatedRunWithMetrics, _ := sw.ee.FetchPodMetrics(run)
	_ = metrics.Timing(metrics.StatusWorkerFetchPodMetrics, time.Since(start), []string{sw.workerId}, 1)

	start = time.Now()
	updatedRun, err := sw.ee.FetchUpdateStatus(updatedRunWithMetrics)
	_ = metrics.Timing(metrics.StatusWorkerFetchUpdateStatus, time.Since(start), []string{sw.workerId}, 1)

	if err == nil {
		subRuns, err := sw.sm.ListRuns(1000, 0, "status", "desc", nil, map[string]string{"PARENT_FLOTILLA_RUN_ID": run.RunID}, state.Engines)
		if err == nil && subRuns.Total > 0 {
			var spawnedRuns state.SpawnedRuns
			for _, subRun := range subRuns.Runs {
				spawnedRuns = append(spawnedRuns, state.SpawnedRun{RunID: subRun.RunID})
			}
			updatedRun.SpawnedRuns = &spawnedRuns
		}
	}
	if err != nil {
		message := fmt.Sprintf("%+v", err)
		_ = sw.log.Log("message", "unable to receive eks runs", "error", message)

		minutesInQueue := time.Now().Sub(*run.QueuedAt).Minutes()
		if strings.Contains(message, "not found") && minutesInQueue > float64(30) {
			stoppedAt := time.Now()
			reason := "Job either timed out or not found on the EKS cluster."
			updatedRun.Status = state.StatusStopped
			updatedRun.FinishedAt = &stoppedAt
			updatedRun.ExitReason = &reason
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}

	} else {
		if run.Status != updatedRun.Status && (updatedRun.PodName == run.PodName) {
			_ = sw.log.Log("message", "updating eks run status", "pod", updatedRun.PodName, "status", updatedRun.Status, "exit_code", updatedRun.ExitCode)

			if sw.exceptionExtractorClient != nil && updatedRun.ExitCode != nil && *(updatedRun.ExitCode) != int64(0) {
				jobUrl := fmt.Sprintf("%s/extract/%s", sw.exceptionExtractorUrl, run.RunID)
				res, err := sw.exceptionExtractorClient.Get(jobUrl)
				if err == nil {
					body, err := ioutil.ReadAll(res.Body)
					if body != nil {
						defer res.Body.Close()
						runExceptions := state.RunExceptions{}
						err = json.Unmarshal(body, &runExceptions)
						if err == nil {
							updatedRun.RunExceptions = &runExceptions
						}
					}
				}
			}
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			if err != nil {
				_ = sw.log.Log("message", "unable to save eks runs", "error", fmt.Sprintf("%+v", err))
			}

			if updatedRun.Status == state.StatusStopped {
				//TODO - move to a separate worker.
				//_ = sw.ee.Terminate(run)
			}
		} else {
			if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
				updatedRun.MaxCpuUsed != run.MaxCpuUsed ||
				updatedRun.Cpu != run.Cpu ||
				updatedRun.Memory != run.Memory ||
				updatedRun.PodEvents != run.PodEvents ||
				updatedRun.SpawnedRuns != run.SpawnedRuns {
				_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
			}
		}
	}
}

func (sw *statusWorker) processEKSRunMetrics(run state.Run) {
	updatedRun, err := sw.ee.FetchPodMetrics(run)
	if err == nil {
		if updatedRun.MaxMemoryUsed != run.MaxMemoryUsed ||
			updatedRun.MaxCpuUsed != run.MaxCpuUsed {
			_, err = sw.sm.UpdateRun(updatedRun.RunID, updatedRun)
		}
	}
}

func (sw *statusWorker) runOnceECS() {
	runReceipt, err := sw.ee.PollStatus()

	if err != nil {
		sw.log.Log("message", "unable to receive status message", "error", fmt.Sprintf("%+v", err))
		return
	}

	// Ensure update is in the env required, otherwise, ack without taking action
	update := runReceipt.Run
	if update != nil {
		//
		// Relies on the reserved env var, FLOTILLA_SERVER_MODE to ensure update
		// belongs to -this- mode of Flotilla
		//
		var serverMode string
		if update.Env != nil {
			for _, kv := range *update.Env {
				if kv.Name == "FLOTILLA_SERVER_MODE" {
					serverMode = kv.Value
				}
			}
		}

		shouldProcess := len(serverMode) > 0 && serverMode == sw.conf.GetString("flotilla_mode")
		if shouldProcess {
			run, err := sw.findRun(update.TaskArn)
			if err != nil {
				sw.log.Log("message", "unable to find run to apply update to", "error", fmt.Sprintf("%+v", err))
				return
			}

			_, err = sw.sm.UpdateRun(run.RunID, *update)
			if err != nil {
				sw.log.Log("message", "error applying status update", "run", run.RunID, "error", fmt.Sprintf("%+v", err))
				return
			}

			// emit status update event
			sw.logStatusUpdate(*update)
		}

		if sw.conf.GetString("flotilla_mode") != "dev" {
			_ = sw.log.Log("message", "Acking status update", "arn", update.TaskArn)
		}
		if err = runReceipt.Done(); err != nil {
			sw.log.Log("message", "Acking status update failed", "arn", update.TaskArn, "error", fmt.Sprintf("%+v", err))
		}
	}
}

func (sw *statusWorker) logStatusUpdate(update state.Run) {
	var err error
	var startedAt, finishedAt time.Time
	var env state.EnvList

	if update.StartedAt != nil {
		startedAt = *update.StartedAt
	}

	if update.FinishedAt != nil {
		finishedAt = *update.FinishedAt
	}

	if update.Env != nil {
		env = *update.Env
	}

	if update.ExitCode != nil {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"exit_code", *update.ExitCode,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	} else {
		err = sw.log.Event("eventClassName", "FlotillaTaskStatus",
			"run_id", update.RunID,
			"task_arn", update.TaskArn,
			"definition_id", update.DefinitionID,
			"alias", update.Alias,
			"image", update.Image,
			"cluster_name", update.ClusterName,
			"status", update.Status,
			"started_at", startedAt,
			"finished_at", finishedAt,
			"instance_id", update.InstanceID,
			"instance_dns_name", update.InstanceDNSName,
			"group_name", update.GroupName,
			"user", update.User,
			"task_type", update.TaskType,
			"env", env,
			"executable_id", update.ExecutableID,
			"executable_type", update.ExecutableType)
	}

	if err != nil {
		sw.log.Log("message", "Failed to emit status event", "run_id", update.RunID, "error", err.Error())
	}
}

func (sw *statusWorker) findRun(taskArn string) (state.Run, error) {
	var engines []string
	if sw.engine != nil {
		engines = []string{*sw.engine}
	} else {
		engines = nil
	}

	runs, err := sw.sm.ListRuns(1, 0, "started_at", "asc", map[string][]string{
		"task_arn": {taskArn},
	}, nil, engines)
	if err != nil {
		return state.Run{}, errors.Wrapf(err, "problem finding run by task arn [%s]", taskArn)
	}
	if runs.Total > 0 && len(runs.Runs) > 0 {
		return runs.Runs[0], nil
	}
	return state.Run{}, errors.Errorf("no run found for [%s]", taskArn)
}
<|endoftext|>


<commit_before>package channel

import (
	"encoding/hex"
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/Hucaru/Valhalla/game/npcchat"

	"github.com/Hucaru/Valhalla/nx"

	"github.com/Hucaru/Valhalla/game"
	"github.com/Hucaru/Valhalla/game/packet"

	"github.com/Hucaru/Valhalla/mnet"
	"github.com/Hucaru/Valhalla/mpacket"
)

func chatSendAll(conn mnet.MConnChannel, reader mpacket.Reader) {
	msg := reader.ReadString(int(reader.ReadInt16()))

	if strings.Index(msg, "/") == 0 && conn.GetAdminLevel() > 0 {
		gmCommand(conn, msg)
	} else {
		player, ok := game.Players[conn]

		if !ok {
			return
		}

		char := player.Char()
		game.Maps[char.MapID].Send(packet.MessageAllChat(char.ID, conn.GetAdminLevel() > 0, msg), player.InstanceID)
	}
}

func chatSlashCommand(conn mnet.MConnChannel, reader mpacket.Reader) {
	cmdType := reader.ReadByte()

	switch cmdType {
	case 5: // FIND
		// length := reader.ReadInt16()
		// name := reader.ReadString(int(length))
	default:
		fmt.Println("Chat command type of", cmdType)
	}
}

func gmCommand(conn mnet.MConnChannel, msg string) {
	ind := strings.Index(msg, "/")
	command := strings.SplitN(msg[ind+1:], " ", -1)

	switch command[0] {
	case "packet":
		if len(command) < 2 {
			return
		}
		packet := string(command[1])
		data, err := hex.DecodeString(packet)

		if err != nil {
			log.Println("Eror in decoding string for gm command packet:", packet)
			break
		}
		log.Println("Sent packet:", hex.EncodeToString(data))
		conn.Send(data)
	case "map":
		var val int
		var err error
		var mapName string

		if len(command) == 2 {
			val, err = strconv.Atoi(command[1])
			mapName = command[1]
		} else if len(command) == 3 {
			val, err = strconv.Atoi(command[2])
			mapName = command[2]
		}

		if err != nil {
			// Check to see if name matches pre-recorded
			switch mapName {
			// Maple island
			case "amherst":
				val = 1010000
			case "southperry":
				val = 60000
			// Victoria island
			case "lith":
				val = 104000000
			case "henesys":
				val = 100000000
			case "kerning":
				val = 103000000
			case "perion":
				val = 102000000
			case "ellinia":
				val = 101000000
			case "sleepy":
				val = 105040300
			case "gm":
				val = 180000000
			// Ossyria
			case "orbis":
				val = 200000000
			case "elnath":
				val = 211000000
			case "ludi":
				val = 220000000
			case "omega":
				val = 221000000
			case "aqua":
				val = 230000000
			// Misc
			case "balrog":
				val = 105090900
			default:
				return
			}
		}

		mapID := int32(val)

		if _, err := nx.GetMap(mapID); err != nil {
			return
		}

		player, ok := game.Players[conn]

		if !ok {
			conn.Send(packet.MessageNotice(err.Error()))
			return
		}

		p, id := game.Maps[player.Char().MapID].GetRandomSpawnPortal()
		player.ChangeMap(mapID, p, id)

	case "notice":
		if len(command) < 2 {
			return
		}
		player, ok := game.Players[conn]

		if !ok {
			conn.Send(packet.MessageNotice("Error in sending notice msg"))
			return
		}

		char := player.Char()
		game.Maps[char.MapID].Send(packet.MessageNotice(strings.Join(command[1:], " ")), player.InstanceID)
	case "kill":
		if len(command) == 1 {
			player, ok := game.Players[conn]

			if !ok {
				conn.Send(packet.MessageNotice("Error in killing player"))
				return
			}

			player.Kill()
		} else {
			if command[1] == "<map>" {
				player, ok := game.Players[conn]

				if !ok {
					conn.Send(packet.MessageNotice("Error in killing players on map"))
					return
				}

				players, err := game.Maps[player.Char().MapID].GetPlayers(player.InstanceID)

				if err != nil {
					return
				}

				for _, v := range players {
					p := game.Players[v]
					p.Kill()
				}

				return
			}

			player, err := game.GetPlayerFromName(command[1])

			if err != nil {
				conn.Send(packet.MessageNotice(err.Error()))
				return
			}

			player.Kill()
		}
	case "revive":
		if len(command) == 1 {
			player, ok := game.Players[conn]

			if !ok {
				conn.Send(packet.MessageNotice("Error in getting player"))
				return
			}

			player.Revive()
		} else {
			if command[1] == "<map>" {
				player, ok := game.Players[conn]

				if !ok {
					conn.Send(packet.MessageNotice("Error in getting player"))
					return
				}

				players, err := game.Maps[player.Char().MapID].GetPlayers(player.InstanceID)

				if err != nil {
					return
				}

				for _, v := range players {
					p := game.Players[v]
					p.Revive()
				}

				return
			}

			player, err := game.GetPlayerFromName(command[1])

			if err != nil {
				conn.Send(packet.MessageNotice(err.Error()))
				return
			}

			player.Revive()
		}
	case "hp":
		if len(command) < 2 {
			return
		}

		player, ok := game.Players[conn]

		if !ok {
			conn.Send(packet.MessageNotice("Error in getting player"))
			return
		}

		if command[1][0] == '+' {
			ammount, err := strconv.Atoi(command[1][1:])

			if err != nil {
				conn.Send(packet.MessageNotice(err.Error()))
			}

			player.GiveHP(int32(ammount))
		} else if command[1][0] == '-' {
			ammount, err := strconv.Atoi(command[1][1:])

			if err != nil {
				conn.Send(packet.MessageNotice(err.Error()))
			}

			player.GiveHP(int32(-ammount))

		} else {
			ammount, err := strconv.Atoi(command[1])

			if err != nil {
				conn.Send(packet.MessageNotice(err.Error()))
			}

			player.SetHP(int32(ammount))
		}
	case "cody":
		if len(command) < 2 {
			return
		}

		npcchat.NewSessionWithOverride(conn, strings.Join(command[1:], " "), 9200000)
		npcchat.Run(conn)

	case "options":
		script := `if state == 1 {
			options = "The following options are available to you:\r\n"
			options += "#dRemember all GM activity is logged\r\n"
			options += "#L0##bSuspicious players#l\r\n"
			options += "#L1#Current map info#l\r\n"
			options += "#L2#Start an event#l\r\n"
			options += "#L3#Spawn mob#l\r\n"

			return SendSelection(options)
		} else if state == 2 {
			if selection == 0 {
				return SendBack("Suspicious player info is not yet recorded")
			} else if selection == 1 {
				return SendBack("Current map info is not yet plumbed")
			} else if selection == 2 {
				return SendBack("Events not coded")
			} else if selection == 3 {
				return SendBack("Why are you trying this?")
			} else {
				return SendBack("Unkown selection")
			}
		}`

		npcchat.NewSessionWithOverride(conn, script, 9010000)
		npcchat.Run(conn)

	case "shop":

	default:
		log.Println("Unkown GM command:", msg)
	}
}
<commit_msg>back to pre instance re-write<commit_after>package channel

import (
	"encoding/hex"
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/Hucaru/Valhalla/game/npcchat"
	"github.com/Hucaru/Valhalla/game/script"

	"github.com/Hucaru/Valhalla/nx"

	"github.com/Hucaru/Valhalla/game"
	"github.com/Hucaru/Valhalla/game/packet"

	"github.com/Hucaru/Valhalla/mnet"
	"github.com/Hucaru/Valhalla/mpacket"
)

func chatSendAll(conn mnet.MConnChannel, reader mpacket.Reader) {
	msg := reader.ReadString(int(reader.ReadInt16()))

	if strings.Index(msg, "/") == 0 && conn.GetAdminLevel() > 0 {
		gmCommand(conn, msg)
	} else {
		player, ok := game.Players[conn]

		if !ok {
			return
		}

		char := player.Char()
		game.Maps[char.MapID].Send(packet.MessageAllChat(char.ID, conn.GetAdminLevel() > 0, msg), player.InstanceID)
	}
}

func chatSlashCommand(conn mnet.MConnChannel, reader mpacket.Reader) {
	cmdType := reader.ReadByte()

	switch cmdType {
	case 5: // FIND
		// length := reader.ReadInt16()
		// name := reader.ReadString(int(length))
	default:
		fmt.Println("Chat command type of", cmdType)
	}
}

func gmCommand(conn mnet.MConnChannel, msg string) {
	ind := strings.Index(msg, "/")
	command := strings.SplitN(msg[ind+1:], " ", -1)

	switch command[0] {
	case "packet":
		if len(command) < 2 {
			return
		}
		packet := string(command[1])
		data, err := hex.DecodeString(packet)

		if err != nil {
			log.Println("Eror in decoding string for gm command packet:", packet)
			break
		}
		log.Println("Sent packet:", hex.EncodeToString(data))
		conn.Send(data)
	case "map":
		var val int
		var err error
		var mapName string

		if len(command) == 2 {
			val, err = strconv.Atoi(command[1])
			mapName = command[1]
		} else if len(command) == 3 {
			val, err = strconv.Atoi(command[2])
			mapName = command[2]
		}

		if err != nil {
			// Check to see if name matches pre-recorded
			switch mapName {
			// Maple island
			case "amherst":
				val = 1010000
			case "southperry":
				val = 60000
			// Victoria island
			case "lith":
				val = 104000000
			case "henesys":
				val = 100000000
			case "kerning":
				val = 103000000
			case "perion":
				val = 102000000
			case "ellinia":
				val = 101000000
			case "sleepy":
				val = 105040300
			case "gm":
				val = 180000000
			// Ossyria
			case "orbis":
				val = 200000000
			case "elnath":
				val = 211000000
			case "ludi":
				val = 220000000
			case "omega":
				val = 221000000
			case "aqua":
				val = 230000000
			// Misc
			case "balrog":
				val = 105090900
			default:
				return
			}
		}

		mapID := int32(val)

		if _, err := nx.GetMap(mapID); err != nil {
			return
		}

		player, ok := game.Players[conn]

		if !ok {
			conn.Send(packet.MessageNotice(err.Error()))
			return
		}

		p, id := game.Maps[player.Char().MapID].GetRandomSpawnPortal()
		player.ChangeMap(mapID, p, id)

	case "notice":
		if len(command) < 2 {
			return
		}
		player, ok := game.Players[conn]

		if !ok {
			conn.Send(packet.MessageNotice("Error in sending notice msg"))
			return
		}

		char := player.Char()
		game.Maps[char.MapID].Send(packet.MessageNotice(strings.Join(command[1:], " ")), player.InstanceID)
	case "kill":
		if len(command) == 1 {
			player, ok := game.Players[conn]

			if !ok {
				conn.Send(packet.MessageNotice("Error in killing player"))
				return
			}

			player.Kill()
		} else {
			if command[1] == "<map>" {
				player, ok := game.Players[conn]

				if !ok {
					conn.Send(packet.MessageNotice("Error in killing players on map"))
					return
				}

				players, err := game.Maps[player.Char().MapID].GetPlayers(player.InstanceID)

				if err != nil {
					return
				}

				for _, v := range players {
					p := game.Players[v]
					p.Kill()
				}

				return
			}

			player, err := game.GetPlayerFromName(command[1])

			if err != nil {
				conn.Send(packet.MessageNotice(err.Error()))
				return
			}

			player.Kill()
		}
	case "revive":
		if len(command) == 1 {
			player, ok := game.Players[conn]

			if !ok {
				conn.Send(packet.MessageNotice("Error in getting player"))
				return
			}

			player.Revive()
		} else {
			if command[1] == "<map>" {
				player, ok := game.Players[conn]

				if !ok {
					conn.Send(packet.MessageNotice("Error in getting player"))
					return
				}

				players, err := game.Maps[player.Char().MapID].GetPlayers(player.InstanceID)

				if err != nil {
					return
				}

				for _, v := range players {
					p := game.Players[v]
					p.Revive()
				}

				return
			}

			player, err := game.GetPlayerFromName(command[1])

			if err != nil {
				conn.Send(packet.MessageNotice(err.Error()))
				return
			}

			player.Revive()
		}
	case "hp":
		if len(command) < 2 {
			return
		}

		player, ok := game.Players[conn]

		if !ok {
			conn.Send(packet.MessageNotice("Error in getting player"))
			return
		}

		if command[1][0] == '+' {
			ammount, err := strconv.Atoi(command[1][1:])

			if err != nil {
				conn.Send(packet.MessageNotice(err.Error()))
			}

			player.GiveHP(int32(ammount))
		} else if command[1][0] == '-' {
			ammount, err := strconv.Atoi(command[1][1:])

			if err != nil {
				conn.Send(packet.MessageNotice(err.Error()))
			}

			player.GiveHP(int32(-ammount))

		} else {
			ammount, err := strconv.Atoi(command[1])

			if err != nil {
				conn.Send(packet.MessageNotice(err.Error()))
			}

			player.SetHP(int32(ammount))
		}
	case "cody":
		if len(command) < 2 {
			return
		}

		npcchat.NewSessionWithOverride(conn, strings.Join(command[1:], " "), 9200000)
		npcchat.Run(conn)

	case "options":
		// script := `if state == 1 {
		// 	options = "The following options are available to you:\r\n"
		// 	options += "#dRemember all GM activity is logged\r\n"
		// 	options += "#L0##bSuspicious players#l\r\n"
		// 	options += "#L1#Current map info#l\r\n"
		// 	options += "#L2#Start an event#l\r\n"
		// 	options += "#L3#Spawn mob#l\r\n"

		// 	return SendSelection(options)
		// } else if state == 2 {
		// 	if selection == 0 {
		// 		return SendBack("Suspicious player info is not yet recorded")
		// 	} else if selection == 1 {
		// 		return SendBack("Current map info is not yet plumbed")
		// 	} else if selection == 2 {
		// 		return SendBack("Events not coded")
		// 	} else if selection == 3 {
		// 		return SendBack("Why are you trying this?")
		// 	} else {
		// 		return SendBack("Unkown selection")
		// 	}
		// }`
		script, err := script.Get("options")

		if err != nil {
			return
		}

		npcchat.NewSessionWithOverride(conn, script, 9010000)
		npcchat.Run(conn)

	case "shop":

	case "createInstance":
	case "changeInstance":
	default:
		log.Println("Unkown GM command:", msg)
	}
}
<|endoftext|>


<commit_before>// Package errors обрабатывает ошибки
package errors

import (
	// Packages
	"log"
	"net/http"
)

// ErrorHandler проверяет ошибки
func ErrorHandler(err error, code int, w http.ResponseWriter) {
	if err != nil {
		log.Print(err) // <= Режим отладки
		if w != nil {
			http.Error(w, http.StatusText(500), 500)
			return
		}
	}
}
<commit_msg>пофиксил модуль ошибок<commit_after>// Package errors обрабатывает ошибки
package errors

import (
	// Packages
	"log"
	"net/http"
)

// ErrorHandler проверяет ошибки
func ErrorHandler(err error, code int, w http.ResponseWriter) {
	if err != nil {
		log.Print(err) // <= Режим отладки
		if w != nil {
			http.Error(w, err.Error(), code)
			return
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 backup_test

import (
	"errors"
	"github.com/jacobsa/comeback/backup"
	"github.com/jacobsa/comeback/backup/mock"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/blob/mock"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/fs/mock"
	"github.com/jacobsa/comeback/repr"
	"github.com/jacobsa/oglemock"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"testing"
)

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

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

func saveBlob(res *[]byte) oglemock.Action {
	f := func(b []byte) (blob.Score, error) {
		*res = b
		return nil, errors.New("foo")
	}

	return oglemock.Invoke(f)
}

type readCloser struct {
	closed bool
}

func (r *readCloser) Read(b []byte) (int, error) {
	panic("Shouldn't be called.")
}

func (r *readCloser) Close() error {
	if r.closed {
		panic("Close called twice.")
	}

	r.closed = true
	return nil
}

func makeFileEntry(name string) *fs.DirectoryEntry {
	return &fs.DirectoryEntry{
		Type: fs.TypeFile,
		Name: name,
	}
}

func makeDirEntry(name string) *fs.DirectoryEntry {
	return &fs.DirectoryEntry{
		Type: fs.TypeDirectory,
		Name: name,
	}
}

type DirectorySaverTest struct {
	blobStore mock_blob.MockStore
	fileSystem mock_fs.MockFileSystem
	fileSaver mock_backup.MockFileSaver
	wrapped mock_backup.MockDirectorySaver

	dirSaver backup.DirectorySaver

	dirpath string
	score blob.Score
	err error
}

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

func (t *DirectorySaverTest) SetUp(i *TestInfo) {
	t.blobStore = mock_blob.NewMockStore(i.MockController, "blobStore")
	t.fileSystem = mock_fs.NewMockFileSystem(i.MockController, "fileSystem")
	t.fileSaver = mock_backup.NewMockFileSaver(i.MockController, "fileSaver")
	t.wrapped = mock_backup.NewMockDirectorySaver(i.MockController, "wrapped")

	t.dirSaver, _ = backup.NewDirectorySaver(
		t.blobStore,
		t.fileSystem,
		t.fileSaver,
		t.wrapped,
	)
}

func (t *DirectorySaverTest) callSaver() {
	t.score, t.err = t.dirSaver.Save(t.dirpath)
}

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

func (t *DirectorySaverTest) CallsReadDir() {
	t.dirpath = "taco"

	// ReadDir
	ExpectCall(t.fileSystem, "ReadDir")("taco").
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.callSaver()
}

func (t *DirectorySaverTest) ReadDirReturnsError() {
	// ReadDir
	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.callSaver()

	ExpectThat(t.err, Error(HasSubstr("Listing")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectorySaverTest) NoEntriesInDirectory() {
	// ReadDir
	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return([]*fs.DirectoryEntry{}, nil))

	// Blob store
	var blob []byte
	ExpectCall(t.blobStore, "Store")(Any()).
		WillOnce(saveBlob(&blob))

	// Call
	t.callSaver()

	AssertNe(nil, blob)
	entries, err := repr.Unmarshal(blob)
	AssertEq(nil, err)
	ExpectThat(entries, ElementsAre())
}

func (t *DirectorySaverTest) CallsFileSystemAndFileSaverForFiles() {
	t.dirpath = "/taco"

	// ReadDir
	entries := []*fs.DirectoryEntry {
		makeFileEntry("burrito"),
		makeFileEntry("enchilada"),
	}

	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(entries, nil))

	// OpenForReading
	file0 := &readCloser{}
	file1 := &readCloser{}

	ExpectCall(t.fileSystem, "OpenForReading")("/taco/burrito").
		WillOnce(oglemock.Return(file0, nil))

	ExpectCall(t.fileSystem, "OpenForReading")("/taco/enchilada").
		WillOnce(oglemock.Return(file1, nil))

	// File saver
	ExpectCall(t.fileSaver, "Save")(file0).
		WillOnce(oglemock.Return([]blob.Score{}, nil))

	ExpectCall(t.fileSaver, "Save")(file1).
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.callSaver()
}

func (t *DirectorySaverTest) FileSystemReturnsErrorForOneFile() {
	// ReadDir
	entries := []*fs.DirectoryEntry {
		makeFileEntry(""),
		makeFileEntry(""),
		makeFileEntry(""),
	}

	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(entries, nil))

	// OpenForReading
	file0 := &readCloser{}

	ExpectCall(t.fileSystem, "OpenForReading")(Any()).
		WillOnce(oglemock.Return(file0, nil)).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// File saver
	ExpectCall(t.fileSaver, "Save")(Any()).
		WillOnce(oglemock.Return([]blob.Score{}, nil))

	// Call
	t.callSaver()

	ExpectThat(t.err, Error(HasSubstr("Opening")))
	ExpectThat(t.err, Error(HasSubstr("taco")))

	ExpectTrue(file0.closed)
}

func (t *DirectorySaverTest) FileSaverReturnsErrorForOneFile() {
	// ReadDir
	entries := []*fs.DirectoryEntry {
		makeFileEntry(""),
		makeFileEntry(""),
		makeFileEntry(""),
	}

	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(entries, nil))

	// OpenForReading
	file0 := &readCloser{}
	file1 := &readCloser{}

	ExpectCall(t.fileSystem, "OpenForReading")(Any()).
		WillOnce(oglemock.Return(file0, nil)).
		WillOnce(oglemock.Return(file1, nil))

	// File saver
	ExpectCall(t.fileSaver, "Save")(Any()).
		WillOnce(oglemock.Return([]blob.Score{}, nil)).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.callSaver()

	ExpectThat(t.err, Error(HasSubstr("taco")))

	ExpectTrue(file0.closed)
	ExpectTrue(file1.closed)
}

func (t *DirectorySaverTest) CallsDirSaverForDirs() {
	t.dirpath = "/taco"

	// ReadDir
	entries := []*fs.DirectoryEntry {
		makeDirEntry("burrito"),
		makeDirEntry("enchilada"),
	}

	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(entries, nil))

	// Wrapped directory saver
	score0 := blob.ComputeScore([]byte(""))

	ExpectCall(t.wrapped, "Save")("/taco/burrito").
		WillOnce(oglemock.Return(score0, nil))

	ExpectCall(t.wrapped, "Save")("/taco/enchilada").
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.callSaver()
}

func (t *DirectorySaverTest) DirSaverReturnsErrorForOneDir() {
	t.dirpath = "/taco"

	// ReadDir
	entries := []*fs.DirectoryEntry {
		makeDirEntry(""),
		makeDirEntry(""),
		makeDirEntry(""),
	}

	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(entries, nil))

	// Wrapped directory saver
	score0 := blob.ComputeScore([]byte(""))

	ExpectCall(t.wrapped, "Save")(Any()).
		WillOnce(oglemock.Return(score0, nil)).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.callSaver()

	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectorySaverTest) OneTypeIsUnsupported() {
	ExpectEq("TODO", "")
}

func (t *DirectorySaverTest) CallsBlobStore() {
	ExpectEq("TODO", "")
}

func (t *DirectorySaverTest) BlobStoreReturnsError() {
	ExpectEq("TODO", "")
}

func (t *DirectorySaverTest) BlobStoreSucceeds() {
	ExpectEq("TODO", "")
}
<commit_msg>DirectorySaverTest.OneTypeIsUnsupported<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 backup_test

import (
	"errors"
	"github.com/jacobsa/comeback/backup"
	"github.com/jacobsa/comeback/backup/mock"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/blob/mock"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/fs/mock"
	"github.com/jacobsa/comeback/repr"
	"github.com/jacobsa/oglemock"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"testing"
)

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

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

func saveBlob(res *[]byte) oglemock.Action {
	f := func(b []byte) (blob.Score, error) {
		*res = b
		return nil, errors.New("foo")
	}

	return oglemock.Invoke(f)
}

type readCloser struct {
	closed bool
}

func (r *readCloser) Read(b []byte) (int, error) {
	panic("Shouldn't be called.")
}

func (r *readCloser) Close() error {
	if r.closed {
		panic("Close called twice.")
	}

	r.closed = true
	return nil
}

func makeFileEntry(name string) *fs.DirectoryEntry {
	return &fs.DirectoryEntry{
		Type: fs.TypeFile,
		Name: name,
	}
}

func makeDirEntry(name string) *fs.DirectoryEntry {
	return &fs.DirectoryEntry{
		Type: fs.TypeDirectory,
		Name: name,
	}
}

type DirectorySaverTest struct {
	blobStore mock_blob.MockStore
	fileSystem mock_fs.MockFileSystem
	fileSaver mock_backup.MockFileSaver
	wrapped mock_backup.MockDirectorySaver

	dirSaver backup.DirectorySaver

	dirpath string
	score blob.Score
	err error
}

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

func (t *DirectorySaverTest) SetUp(i *TestInfo) {
	t.blobStore = mock_blob.NewMockStore(i.MockController, "blobStore")
	t.fileSystem = mock_fs.NewMockFileSystem(i.MockController, "fileSystem")
	t.fileSaver = mock_backup.NewMockFileSaver(i.MockController, "fileSaver")
	t.wrapped = mock_backup.NewMockDirectorySaver(i.MockController, "wrapped")

	t.dirSaver, _ = backup.NewDirectorySaver(
		t.blobStore,
		t.fileSystem,
		t.fileSaver,
		t.wrapped,
	)
}

func (t *DirectorySaverTest) callSaver() {
	t.score, t.err = t.dirSaver.Save(t.dirpath)
}

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

func (t *DirectorySaverTest) CallsReadDir() {
	t.dirpath = "taco"

	// ReadDir
	ExpectCall(t.fileSystem, "ReadDir")("taco").
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.callSaver()
}

func (t *DirectorySaverTest) ReadDirReturnsError() {
	// ReadDir
	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.callSaver()

	ExpectThat(t.err, Error(HasSubstr("Listing")))
	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectorySaverTest) NoEntriesInDirectory() {
	// ReadDir
	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return([]*fs.DirectoryEntry{}, nil))

	// Blob store
	var blob []byte
	ExpectCall(t.blobStore, "Store")(Any()).
		WillOnce(saveBlob(&blob))

	// Call
	t.callSaver()

	AssertNe(nil, blob)
	entries, err := repr.Unmarshal(blob)
	AssertEq(nil, err)
	ExpectThat(entries, ElementsAre())
}

func (t *DirectorySaverTest) CallsFileSystemAndFileSaverForFiles() {
	t.dirpath = "/taco"

	// ReadDir
	entries := []*fs.DirectoryEntry {
		makeFileEntry("burrito"),
		makeFileEntry("enchilada"),
	}

	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(entries, nil))

	// OpenForReading
	file0 := &readCloser{}
	file1 := &readCloser{}

	ExpectCall(t.fileSystem, "OpenForReading")("/taco/burrito").
		WillOnce(oglemock.Return(file0, nil))

	ExpectCall(t.fileSystem, "OpenForReading")("/taco/enchilada").
		WillOnce(oglemock.Return(file1, nil))

	// File saver
	ExpectCall(t.fileSaver, "Save")(file0).
		WillOnce(oglemock.Return([]blob.Score{}, nil))

	ExpectCall(t.fileSaver, "Save")(file1).
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.callSaver()
}

func (t *DirectorySaverTest) FileSystemReturnsErrorForOneFile() {
	// ReadDir
	entries := []*fs.DirectoryEntry {
		makeFileEntry(""),
		makeFileEntry(""),
		makeFileEntry(""),
	}

	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(entries, nil))

	// OpenForReading
	file0 := &readCloser{}

	ExpectCall(t.fileSystem, "OpenForReading")(Any()).
		WillOnce(oglemock.Return(file0, nil)).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// File saver
	ExpectCall(t.fileSaver, "Save")(Any()).
		WillOnce(oglemock.Return([]blob.Score{}, nil))

	// Call
	t.callSaver()

	ExpectThat(t.err, Error(HasSubstr("Opening")))
	ExpectThat(t.err, Error(HasSubstr("taco")))

	ExpectTrue(file0.closed)
}

func (t *DirectorySaverTest) FileSaverReturnsErrorForOneFile() {
	// ReadDir
	entries := []*fs.DirectoryEntry {
		makeFileEntry(""),
		makeFileEntry(""),
		makeFileEntry(""),
	}

	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(entries, nil))

	// OpenForReading
	file0 := &readCloser{}
	file1 := &readCloser{}

	ExpectCall(t.fileSystem, "OpenForReading")(Any()).
		WillOnce(oglemock.Return(file0, nil)).
		WillOnce(oglemock.Return(file1, nil))

	// File saver
	ExpectCall(t.fileSaver, "Save")(Any()).
		WillOnce(oglemock.Return([]blob.Score{}, nil)).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.callSaver()

	ExpectThat(t.err, Error(HasSubstr("taco")))

	ExpectTrue(file0.closed)
	ExpectTrue(file1.closed)
}

func (t *DirectorySaverTest) CallsDirSaverForDirs() {
	t.dirpath = "/taco"

	// ReadDir
	entries := []*fs.DirectoryEntry {
		makeDirEntry("burrito"),
		makeDirEntry("enchilada"),
	}

	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(entries, nil))

	// Wrapped directory saver
	score0 := blob.ComputeScore([]byte(""))

	ExpectCall(t.wrapped, "Save")("/taco/burrito").
		WillOnce(oglemock.Return(score0, nil))

	ExpectCall(t.wrapped, "Save")("/taco/enchilada").
		WillOnce(oglemock.Return(nil, errors.New("")))

	// Call
	t.callSaver()
}

func (t *DirectorySaverTest) DirSaverReturnsErrorForOneDir() {
	t.dirpath = "/taco"

	// ReadDir
	entries := []*fs.DirectoryEntry {
		makeDirEntry(""),
		makeDirEntry(""),
		makeDirEntry(""),
	}

	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(entries, nil))

	// Wrapped directory saver
	score0 := blob.ComputeScore([]byte(""))

	ExpectCall(t.wrapped, "Save")(Any()).
		WillOnce(oglemock.Return(score0, nil)).
		WillOnce(oglemock.Return(nil, errors.New("taco")))

	// Call
	t.callSaver()

	ExpectThat(t.err, Error(HasSubstr("taco")))
}

func (t *DirectorySaverTest) OneTypeIsUnsupported() {
	t.dirpath = "/taco"

	// ReadDir
	entries := []*fs.DirectoryEntry {
		makeDirEntry(""),
		makeDirEntry(""),
		makeDirEntry(""),
	}

	entries[1].Type = fs.TypeSymlink

	ExpectCall(t.fileSystem, "ReadDir")(Any()).
		WillOnce(oglemock.Return(entries, nil))

	// Wrapped directory saver
	score0 := blob.ComputeScore([]byte(""))

	ExpectCall(t.wrapped, "Save")(Any()).
		WillOnce(oglemock.Return(score0, nil))

	// Call
	t.callSaver()

	ExpectThat(t.err, Error(HasSubstr("Unhandled")))
	ExpectThat(t.err, Error(HasSubstr("type")))
}

func (t *DirectorySaverTest) CallsBlobStore() {
	ExpectEq("TODO", "")
}

func (t *DirectorySaverTest) BlobStoreReturnsError() {
	ExpectEq("TODO", "")
}

func (t *DirectorySaverTest) BlobStoreSucceeds() {
	ExpectEq("TODO", "")
}
<|endoftext|>


<commit_before>package amazonebs

import (
	"fmt"
	"github.com/mitchellh/goamz/ec2"
	"github.com/mitchellh/packer/packer"
	"log"
	"strings"
)

type artifact struct {
	// A map of regions to AMI IDs.
	amis map[string]string

	// EC2 connection for performing API stuff.
	conn *ec2.EC2
}

func (*artifact) BuilderId() string {
	return BuilderId
}

func (*artifact) Files() []string {
	// We have no files
	return nil
}

func (a *artifact) Id() string {
	parts := make([]string, 0, len(a.amis))
	for region, amiId := range a.amis {
		parts = append(parts, fmt.Sprintf("%s:%s", region, amiId))
	}

	return strings.Join(parts, ",")
}

func (a *artifact) String() string {
	amiStrings := make([]string, 0, len(a.amis))
	for region, id := range a.amis {
		single := fmt.Sprintf("%s: %s", region, id)
		amiStrings = append(amiStrings, single)
	}

	return fmt.Sprintf("AMIs were created:\n\n%s", strings.Join(amiStrings, "\n"))
}

func (a *artifact) Destroy() error {
	errors := make([]error, 0)

	for _, imageId := range a.amis {
		log.Printf("Degistering image ID: %s", imageId)
		if _, err := a.conn.DeregisterImage(imageId); err != nil {
			errors = append(errors, err)
		}

		// TODO(mitchellh): Delete the snapshots associated with an AMI too
	}

	if len(errors) > 0 {
		if len(errors) == 1 {
			return errors[0]
		} else {
			return &packer.MultiError{errors}
		}
	}

	return nil
}
<commit_msg>Fix a typo.<commit_after>package amazonebs

import (
	"fmt"
	"github.com/mitchellh/goamz/ec2"
	"github.com/mitchellh/packer/packer"
	"log"
	"strings"
)

type artifact struct {
	// A map of regions to AMI IDs.
	amis map[string]string

	// EC2 connection for performing API stuff.
	conn *ec2.EC2
}

func (*artifact) BuilderId() string {
	return BuilderId
}

func (*artifact) Files() []string {
	// We have no files
	return nil
}

func (a *artifact) Id() string {
	parts := make([]string, 0, len(a.amis))
	for region, amiId := range a.amis {
		parts = append(parts, fmt.Sprintf("%s:%s", region, amiId))
	}

	return strings.Join(parts, ",")
}

func (a *artifact) String() string {
	amiStrings := make([]string, 0, len(a.amis))
	for region, id := range a.amis {
		single := fmt.Sprintf("%s: %s", region, id)
		amiStrings = append(amiStrings, single)
	}

	return fmt.Sprintf("AMIs were created:\n\n%s", strings.Join(amiStrings, "\n"))
}

func (a *artifact) Destroy() error {
	errors := make([]error, 0)

	for _, imageId := range a.amis {
		log.Printf("Deregistering image ID: %s", imageId)
		if _, err := a.conn.DeregisterImage(imageId); err != nil {
			errors = append(errors, err)
		}

		// TODO(mitchellh): Delete the snapshots associated with an AMI too
	}

	if len(errors) > 0 {
		if len(errors) == 1 {
			return errors[0]
		} else {
			return &packer.MultiError{errors}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package library

import (
	"fmt"
	"os"
	"path"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/docker-library/go-dockerlibrary/manifest"
	"github.com/docker/distribution/reference"
	git "gopkg.in/src-d/go-git.v4"
	"gopkg.in/src-d/go-git.v4/plumbing"
)

// LibraryRepo is the Git URL to the official images repository.
var LibraryRepo = "https://github.com/docker-library/official-images.git"

// LibraryResolver attempts to resolve the source of a library Dockerfile.
type LibraryResolver struct {
	libraryPath   string
	repositoryDir string
}

// GetLibrarySource clones the Dockerfile source for a library image.
func (r *LibraryResolver) GetLibrarySource(ref reference.NamedTagged) (string, error) {
	le := log.WithField("ref", ref.String())
	refName := ref.Name()
	if strings.HasPrefix(refName, "library/") {
		refName = refName[len("library/"):]
	}
	le.Debug("Consulting manifest")
	_, _, man, err := manifest.Fetch(r.libraryPath, refName)
	if err != nil {
		return "", err
	}
	tag := man.GetTag(ref.Tag())
	if tag == nil {
		return "", fmt.Errorf("Cannot find tag %s in library repo %s", ref.Tag(), refName)
	}
	le.Debugf("Got tag:\n%s", tag.String())
	repoPath := path.Join(r.repositoryDir, "library-"+refName)
	gle := log.WithField("repo", tag.GitRepo)
	var repo *git.Repository
	if _, err := os.Stat(repoPath); err == nil {
		repo, _ = git.PlainOpen(repoPath)
	}
	if repo == nil {
		os.RemoveAll(repoPath)
		gle.WithField("path", repoPath).Debug("Cloning")
		repo, err = git.PlainClone(repoPath, false, &git.CloneOptions{
			URL:               tag.GitRepo,
			RecurseSubmodules: git.DefaultSubmoduleRecursionDepth,
			Progress:          os.Stdout,
		})
		if err != nil {
			return "", err
		}
	}
	wt, err := repo.Worktree()
	if err != nil {
		return "", err
	}
	if tag.GitCommit != "" {
		le.WithField("ref", tag.GitCommit).Debug("Checking-out")
		err = wt.Checkout(&git.CheckoutOptions{
			Hash: plumbing.NewHash(tag.GitCommit),
		})
		if err != nil {
			status := "unknown"
			stat, err := wt.Status()
			if err != nil {
				status = stat.String()
			}
			le.WithError(err).Warnf("Unable to checkout commit, using:\n%s", status)
		}
	}
	dockerfileDir := path.Join(repoPath, tag.Directory)
	return dockerfileDir, nil
}

// NewLibraryResolver builds a library resolver.
func NewLibraryResolver(libraryPath string, repositoryDir string) *LibraryResolver {
	return &LibraryResolver{libraryPath: libraryPath, repositoryDir: repositoryDir}
}

// BuildLibraryResolver sets up a library resolver given a temporary cache directory.
func BuildLibraryResolver(cacheDir string) (*LibraryResolver, error) {
	repoDir := path.Join(cacheDir, "official-images")
	log.WithField("url", LibraryRepo).Debug("Cloning")
	var repo *git.Repository
	if st, err := os.Stat(repoDir); err == nil && st.IsDir() {
		r, err := git.PlainOpen(repoDir)
		if err == nil {
			repo = r
			r.Fetch(&git.FetchOptions{
				Progress:   os.Stdout,
				RemoteName: "origin",
			})
			wt, err := r.Worktree()
			if err == nil {
				wt.Checkout(&git.CheckoutOptions{
					Branch: plumbing.ReferenceName("master"),
					Force:  true,
				})
			}
		}
	} else {
		os.RemoveAll(repoDir)
	}
	if repo == nil {
		_, err := git.PlainClone(repoDir, false, &git.CloneOptions{
			Depth:    1,
			Progress: os.Stdout,
			URL:      LibraryRepo,
		})
		if err != nil {
			return nil, err
		}
		// repo = r
	}
	return NewLibraryResolver(path.Join(repoDir, "library"), cacheDir), nil
}
<commit_msg>fix: scratchbuild - trim docker.io prefix<commit_after>package library

import (
	"fmt"
	"os"
	"path"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/docker-library/go-dockerlibrary/manifest"
	"github.com/docker/distribution/reference"
	git "gopkg.in/src-d/go-git.v4"
	"gopkg.in/src-d/go-git.v4/plumbing"
)

// LibraryRepo is the Git URL to the official images repository.
var LibraryRepo = "https://github.com/docker-library/official-images.git"

// LibraryResolver attempts to resolve the source of a library Dockerfile.
type LibraryResolver struct {
	libraryPath   string
	repositoryDir string
}

// GetLibrarySource clones the Dockerfile source for a library image.
func (r *LibraryResolver) GetLibrarySource(ref reference.NamedTagged) (string, error) {
	le := log.WithField("ref", ref.String())
	refName := ref.Name()
	if strings.HasPrefix(refName, "docker.io/") {
		refName = refName[len("docker.io/"):]
	}
	if strings.HasPrefix(refName, "library/") {
		refName = refName[len("library/"):]
	}
	le.Debug("Consulting manifest")
	_, _, man, err := manifest.Fetch(r.libraryPath, refName)
	if err != nil {
		return "", err
	}
	tag := man.GetTag(ref.Tag())
	if tag == nil {
		return "", fmt.Errorf("Cannot find tag %s in library repo %s", ref.Tag(), refName)
	}
	le.Debugf("Got tag:\n%s", tag.String())
	repoPath := path.Join(r.repositoryDir, "library-"+refName)
	gle := log.WithField("repo", tag.GitRepo)
	var repo *git.Repository
	if _, err := os.Stat(repoPath); err == nil {
		repo, _ = git.PlainOpen(repoPath)
	}
	if repo == nil {
		os.RemoveAll(repoPath)
		gle.WithField("path", repoPath).Debug("Cloning")
		repo, err = git.PlainClone(repoPath, false, &git.CloneOptions{
			URL:               tag.GitRepo,
			RecurseSubmodules: git.DefaultSubmoduleRecursionDepth,
			Progress:          os.Stdout,
		})
		if err != nil {
			return "", err
		}
	}
	wt, err := repo.Worktree()
	if err != nil {
		return "", err
	}
	if tag.GitCommit != "" {
		le.WithField("ref", tag.GitCommit).Debug("Checking-out")
		err = wt.Checkout(&git.CheckoutOptions{
			Hash: plumbing.NewHash(tag.GitCommit),
		})
		if err != nil {
			status := "unknown"
			stat, err := wt.Status()
			if err != nil {
				status = stat.String()
			}
			le.WithError(err).Warnf("Unable to checkout commit, using:\n%s", status)
		}
	}
	dockerfileDir := path.Join(repoPath, tag.Directory)
	return dockerfileDir, nil
}

// NewLibraryResolver builds a library resolver.
func NewLibraryResolver(libraryPath string, repositoryDir string) *LibraryResolver {
	return &LibraryResolver{libraryPath: libraryPath, repositoryDir: repositoryDir}
}

// BuildLibraryResolver sets up a library resolver given a temporary cache directory.
func BuildLibraryResolver(cacheDir string) (*LibraryResolver, error) {
	repoDir := path.Join(cacheDir, "official-images")
	log.WithField("url", LibraryRepo).Debug("Cloning")
	var repo *git.Repository
	if st, err := os.Stat(repoDir); err == nil && st.IsDir() {
		r, err := git.PlainOpen(repoDir)
		if err == nil {
			repo = r
			r.Fetch(&git.FetchOptions{
				Progress:   os.Stdout,
				RemoteName: "origin",
			})
			wt, err := r.Worktree()
			if err == nil {
				wt.Checkout(&git.CheckoutOptions{
					Branch: plumbing.ReferenceName("master"),
					Force:  true,
				})
			}
		}
	} else {
		os.RemoveAll(repoDir)
	}
	if repo == nil {
		_, err := git.PlainClone(repoDir, false, &git.CloneOptions{
			Depth:    1,
			Progress: os.Stdout,
			URL:      LibraryRepo,
		})
		if err != nil {
			return nil, err
		}
		// repo = r
	}
	return NewLibraryResolver(path.Join(repoDir, "library"), cacheDir), nil
}
<|endoftext|>


<commit_before>package repository

import (
	"github.com/lib/pq"
	"github.com/syou6162/go-active-learning/lib/model"
)

func (r *repository) UpdateOrCreateReferringTweets(e *model.Example) error {
	if e.ReferringTweets == nil || len(*e.ReferringTweets) == 0 {
		return nil
	}

	tmp, err := r.FindExampleByUlr(e.Url)
	if err != nil {
		return err
	}
	id := tmp.Id

	for _, t := range *e.ReferringTweets {
		t.ExampleId = id
		if _, err = r.db.NamedExec(`
INSERT INTO tweet
( example_id,  created_at,  id_str,  full_text,  favorite_count,  retweet_count,  lang,  screen_name,  name,  profile_image_url,  label,  score)
VALUES
(:example_id, :created_at, :id_str, :full_text, :favorite_count, :retweet_count, :lang, :screen_name, :name, :profile_image_url, :label, :score)
ON CONFLICT (example_id, id_str)
DO UPDATE SET
favorite_count = :favorite_count,  retweet_count = :retweet_count, label = :label
WHERE
EXCLUDED.label != 0 AND tweet.label != EXCLUDED.label
;`, t); err != nil {
			return err
		}
	}
	return nil
}

func (r *repository) UpdateTweetLabel(exampleId int, idStr string, label model.LabelType) error {
	if _, err := r.db.Exec(`UPDATE tweet SET label = $1 WHERE example_id = $2 AND id_str = $3;`, label, exampleId, idStr); err != nil {
		return err
	}
	return nil
}

func (r *repository) SearchReferringTweetsList(examples model.Examples) (map[int]model.ReferringTweets, error) {
	referringTweetsByExampleId := make(map[int]model.ReferringTweets)

	referringTweets := model.ReferringTweets{}
	exampleIds := make([]int, 0)
	for _, e := range examples {
		exampleIds = append(exampleIds, e.Id)
	}

	query := `SELECT * FROM tweet WHERE example_id = ANY($1) AND label != -1 AND score > -1.0 AND (lang = 'en' OR lang = 'ja') ORDER BY favorite_count DESC;`
	err := r.db.Select(&referringTweets, query, pq.Array(exampleIds))
	if err != nil {
		return referringTweetsByExampleId, err
	}

	for _, t := range referringTweets {
		referringTweetsByExampleId[t.ExampleId] = append(referringTweetsByExampleId[t.ExampleId], t)
	}
	return referringTweetsByExampleId, nil
}

func (r *repository) SearchReferringTweets(limit int) (model.ReferringTweets, error) {
	referringTweets := model.ReferringTweets{}
	query := `SELECT * FROM tweet lang = 'en' OR lang = 'ja' ORDER BY created_at DESC LIMIT $1;`
	err := r.db.Select(&referringTweets, query, limit)
	if err != nil {
		return referringTweets, err
	}
	return referringTweets, nil
}

func (r *repository) searchReferringTweetsByLabel(label model.LabelType, limit int) (model.ReferringTweets, error) {
	referringTweets := model.ReferringTweets{}
	query := `SELECT * FROM tweet WHERE label = $1 AND (lang = 'en' OR lang = 'ja') ORDER BY created_at DESC LIMIT $2;`
	err := r.db.Select(&referringTweets, query, label, limit)
	if err != nil {
		return referringTweets, err
	}
	return referringTweets, nil
}

func (r *repository) SearchPositiveReferringTweets(limit int) (model.ReferringTweets, error) {
	return r.searchReferringTweetsByLabel(model.POSITIVE, limit)
}

func (r *repository) SearchNegativeReferringTweets(limit int) (model.ReferringTweets, error) {
	return r.searchReferringTweetsByLabel(model.NEGATIVE, limit)
}

func (r *repository) SearchUnlabeledReferringTweets(limit int) (model.ReferringTweets, error) {
	return r.searchReferringTweetsByLabel(model.UNLABELED, limit)
}

func (r *repository) FindReferringTweets(e *model.Example) (model.ReferringTweets, error) {
	referringTweets := model.ReferringTweets{}

	query := `SELECT * FROM tweet WHERE example_id = $1 AND label != -1 AND score > -1.0 AND (lang = 'en' OR lang = 'ja') ORDER BY favorite_count DESC;`
	err := r.db.Select(&referringTweets, query, e.Id)
	if err != nil {
		return referringTweets, err
	}
	return referringTweets, nil
}
<commit_msg>where抜けてる<commit_after>package repository

import (
	"github.com/lib/pq"
	"github.com/syou6162/go-active-learning/lib/model"
)

func (r *repository) UpdateOrCreateReferringTweets(e *model.Example) error {
	if e.ReferringTweets == nil || len(*e.ReferringTweets) == 0 {
		return nil
	}

	tmp, err := r.FindExampleByUlr(e.Url)
	if err != nil {
		return err
	}
	id := tmp.Id

	for _, t := range *e.ReferringTweets {
		t.ExampleId = id
		if _, err = r.db.NamedExec(`
INSERT INTO tweet
( example_id,  created_at,  id_str,  full_text,  favorite_count,  retweet_count,  lang,  screen_name,  name,  profile_image_url,  label,  score)
VALUES
(:example_id, :created_at, :id_str, :full_text, :favorite_count, :retweet_count, :lang, :screen_name, :name, :profile_image_url, :label, :score)
ON CONFLICT (example_id, id_str)
DO UPDATE SET
favorite_count = :favorite_count,  retweet_count = :retweet_count, label = :label
WHERE
EXCLUDED.label != 0 AND tweet.label != EXCLUDED.label
;`, t); err != nil {
			return err
		}
	}
	return nil
}

func (r *repository) UpdateTweetLabel(exampleId int, idStr string, label model.LabelType) error {
	if _, err := r.db.Exec(`UPDATE tweet SET label = $1 WHERE example_id = $2 AND id_str = $3;`, label, exampleId, idStr); err != nil {
		return err
	}
	return nil
}

func (r *repository) SearchReferringTweetsList(examples model.Examples) (map[int]model.ReferringTweets, error) {
	referringTweetsByExampleId := make(map[int]model.ReferringTweets)

	referringTweets := model.ReferringTweets{}
	exampleIds := make([]int, 0)
	for _, e := range examples {
		exampleIds = append(exampleIds, e.Id)
	}

	query := `SELECT * FROM tweet WHERE example_id = ANY($1) AND label != -1 AND score > -1.0 AND (lang = 'en' OR lang = 'ja') ORDER BY favorite_count DESC;`
	err := r.db.Select(&referringTweets, query, pq.Array(exampleIds))
	if err != nil {
		return referringTweetsByExampleId, err
	}

	for _, t := range referringTweets {
		referringTweetsByExampleId[t.ExampleId] = append(referringTweetsByExampleId[t.ExampleId], t)
	}
	return referringTweetsByExampleId, nil
}

func (r *repository) SearchReferringTweets(limit int) (model.ReferringTweets, error) {
	referringTweets := model.ReferringTweets{}
	query := `SELECT * FROM WHERE tweet lang = 'en' OR lang = 'ja' ORDER BY created_at DESC LIMIT $1;`
	err := r.db.Select(&referringTweets, query, limit)
	if err != nil {
		return referringTweets, err
	}
	return referringTweets, nil
}

func (r *repository) searchReferringTweetsByLabel(label model.LabelType, limit int) (model.ReferringTweets, error) {
	referringTweets := model.ReferringTweets{}
	query := `SELECT * FROM tweet WHERE label = $1 AND (lang = 'en' OR lang = 'ja') ORDER BY created_at DESC LIMIT $2;`
	err := r.db.Select(&referringTweets, query, label, limit)
	if err != nil {
		return referringTweets, err
	}
	return referringTweets, nil
}

func (r *repository) SearchPositiveReferringTweets(limit int) (model.ReferringTweets, error) {
	return r.searchReferringTweetsByLabel(model.POSITIVE, limit)
}

func (r *repository) SearchNegativeReferringTweets(limit int) (model.ReferringTweets, error) {
	return r.searchReferringTweetsByLabel(model.NEGATIVE, limit)
}

func (r *repository) SearchUnlabeledReferringTweets(limit int) (model.ReferringTweets, error) {
	return r.searchReferringTweetsByLabel(model.UNLABELED, limit)
}

func (r *repository) FindReferringTweets(e *model.Example) (model.ReferringTweets, error) {
	referringTweets := model.ReferringTweets{}

	query := `SELECT * FROM tweet WHERE example_id = $1 AND label != -1 AND score > -1.0 AND (lang = 'en' OR lang = 'ja') ORDER BY favorite_count DESC;`
	err := r.db.Select(&referringTweets, query, e.Id)
	if err != nil {
		return referringTweets, err
	}
	return referringTweets, nil
}
<|endoftext|>


<commit_before>package base

import (
	"encoding/json"
	"fmt"
	"sync/atomic"

	sgbucket "github.com/couchbase/sg-bucket"
)

// A wrapper around a Bucket that tracks bucket usage statistics as basic read/write counts.  Doesn't break
// down by operation type, to better identify counts for bulk operations
type StatsBucket struct {
	bucket             Bucket
	docsRead           uint64
	docsWritten        uint64
	bytesRead          uint64
	bytesWritten       uint64
	unknownSizeRead    uint64
	unknownSizeWritten uint64
}

type StatsBucketStats struct {
	DocsRead           uint64
	DocsWritten        uint64
	BytesRead          uint64
	BytesWritten       uint64
	UnknownSizeRead    uint64
	UnknownSizeWritten uint64
}

func (sbs *StatsBucketStats) String() string {
	return fmt.Sprintf("\nDocs Read:          %12d \nDocs Written:       %12d \nBytes Read:         %12d \nBytes Written:      %12d \nUnknown Size Reads: %12d \nUnknown Size Writes:%12d",
		sbs.DocsRead, sbs.DocsWritten, sbs.BytesRead, sbs.BytesWritten, sbs.UnknownSizeRead, sbs.UnknownSizeWritten)
}

func (sbs *StatsBucketStats) PerIteration(iterationCount uint64) string {
	return fmt.Sprintf("\nDocs Read:          %12d \nDocs Written:       %12d \nBytes Read:         %12d \nBytes Written:      %12d \nUnknown Size Reads: %12d \nUnknown Size Writes:%12d",
		sbs.DocsRead/iterationCount,
		sbs.DocsWritten/iterationCount,
		sbs.BytesRead/iterationCount,
		sbs.BytesWritten/iterationCount,
		sbs.UnknownSizeRead/iterationCount,
		sbs.UnknownSizeWritten/iterationCount)
}

func NewStatsBucket(bucket Bucket) *StatsBucket {
	return &StatsBucket{
		bucket: bucket,
	}
}

func (b *StatsBucket) docRead(count, bytesRead int) {
	atomic.AddUint64(&b.docsRead, uint64(count))
	if bytesRead == -1 {
		atomic.AddUint64(&b.unknownSizeRead, uint64(1))
	} else {
		atomic.AddUint64(&b.bytesRead, uint64(bytesRead))
	}
}

func (b *StatsBucket) docWrite(count, bytesWritten int) {
	atomic.AddUint64(&b.docsWritten, uint64(count))
	if bytesWritten == -1 {
		atomic.AddUint64(&b.unknownSizeWritten, uint64(1))
	} else {
		atomic.AddUint64(&b.bytesWritten, uint64(bytesWritten))

	}
}

func (b *StatsBucket) GetStats() StatsBucketStats {
	return StatsBucketStats{
		DocsRead:           atomic.LoadUint64(&b.docsRead),
		DocsWritten:        atomic.LoadUint64(&b.docsWritten),
		BytesRead:          atomic.LoadUint64(&b.bytesRead),
		BytesWritten:       atomic.LoadUint64(&b.bytesWritten),
		UnknownSizeRead:    atomic.LoadUint64(&b.unknownSizeRead),
		UnknownSizeWritten: atomic.LoadUint64(&b.unknownSizeWritten),
	}
}

func (b *StatsBucket) GetName() string {
	return b.bucket.GetName()
}
func (b *StatsBucket) Get(k string, rv interface{}) (uint64, error) {

	cas, err := b.bucket.Get(k, rv)
	if vBytes, ok := rv.([]byte); ok {
		defer b.docRead(1, len(vBytes))
	} else if marshalledJSON, marshalErr := json.Marshal(rv); marshalErr == nil {
		defer b.docRead(1, len(marshalledJSON))
	} else {
		defer b.docRead(1, -1)
	}
	return cas, err
}
func (b *StatsBucket) GetRaw(k string) (v []byte, cas uint64, err error) {
	v, cas, err = b.bucket.GetRaw(k)
	b.docRead(1, len(v))
	return v, cas, err
}
func (b *StatsBucket) GetAndTouchRaw(k string, exp int) (v []byte, cas uint64, err error) {
	v, cas, err = b.bucket.GetAndTouchRaw(k, exp)
	b.docRead(1, len(v))
	return v, cas, err
}
func (b *StatsBucket) GetBulkRaw(keys []string) (map[string][]byte, error) {
	results, err := b.bucket.GetBulkRaw(keys)
	for _, value := range results {
		b.docRead(1, len(value))
	}
	return results, err
}
func (b *StatsBucket) Add(k string, exp int, v interface{}) (added bool, err error) {
	if vBytes, ok := v.([]byte); ok {
		defer b.docWrite(1, len(vBytes))
	} else {
		defer b.docWrite(1, -1)
	}
	return b.bucket.Add(k, exp, v)
}
func (b *StatsBucket) AddRaw(k string, exp int, v []byte) (added bool, err error) {
	defer b.docWrite(1, len(v))
	return b.bucket.AddRaw(k, exp, v)
}
func (b *StatsBucket) Append(k string, data []byte) error {
	defer b.docWrite(1, len(data))
	return b.bucket.Append(k, data)
}
func (b *StatsBucket) Set(k string, exp int, v interface{}) error {
	if vBytes, ok := v.([]byte); ok {
		defer b.docWrite(1, len(vBytes))
	} else {
		defer b.docWrite(1, -1)
	}
	return b.bucket.Set(k, exp, v)
}
func (b *StatsBucket) SetRaw(k string, exp int, v []byte) error {
	defer b.docWrite(1, len(v))
	return b.bucket.SetRaw(k, exp, v)
}
func (b *StatsBucket) Delete(k string) error {
	return b.bucket.Delete(k)
}
func (b *StatsBucket) Write(k string, flags int, exp int, v interface{}, opt sgbucket.WriteOptions) error {
	if vBytes, ok := v.([]byte); ok {
		defer b.docWrite(1, len(vBytes))
	} else {
		defer b.docWrite(1, -1)
	}
	return b.bucket.Write(k, flags, exp, v, opt)
}
func (b *StatsBucket) WriteCas(k string, flags int, exp int, cas uint64, v interface{}, opt sgbucket.WriteOptions) (uint64, error) {
	if vBytes, ok := v.([]byte); ok {
		defer b.docWrite(1, len(vBytes))
	} else {
		defer b.docWrite(1, -1)
	}
	return b.bucket.WriteCas(k, flags, exp, cas, v, opt)
}
func (b *StatsBucket) Update(k string, exp int, callback sgbucket.UpdateFunc) (err error) {
	defer b.docWrite(1, -1)
	return b.bucket.Update(k, exp, callback)
}
func (b *StatsBucket) WriteUpdate(k string, exp int, callback sgbucket.WriteUpdateFunc) (err error) {
	defer b.docWrite(1, -1)
	return b.bucket.WriteUpdate(k, exp, callback)
}
func (b *StatsBucket) Incr(k string, amt, def uint64, exp int) (uint64, error) {
	return b.bucket.Incr(k, amt, def, exp)
}
func (b *StatsBucket) GetDDoc(docname string, value interface{}) error {
	return b.bucket.GetDDoc(docname, value)
}
func (b *StatsBucket) PutDDoc(docname string, value interface{}) error {
	return b.bucket.PutDDoc(docname, value)
}
func (b *StatsBucket) DeleteDDoc(docname string) error {
	return b.bucket.DeleteDDoc(docname)
}
func (b *StatsBucket) View(ddoc, name string, params map[string]interface{}) (sgbucket.ViewResult, error) {
	return b.bucket.View(ddoc, name, params)
}

func (b *StatsBucket) ViewCustom(ddoc, name string, params map[string]interface{}, vres interface{}) error {
	return b.bucket.ViewCustom(ddoc, name, params, vres)
}

func (b *StatsBucket) SetBulk(entries []*sgbucket.BulkSetEntry) (err error) {
	defer b.docWrite(len(entries), 0)
	return b.bucket.SetBulk(entries)
}

func (b *StatsBucket) Refresh() error {
	return b.bucket.Refresh()
}

func (b *StatsBucket) StartTapFeed(args sgbucket.TapArguments) (sgbucket.TapFeed, error) {
	return b.bucket.StartTapFeed(args)
}
func (b *StatsBucket) Close() {
	b.bucket.Close()
}
func (b *StatsBucket) Dump() {
	b.bucket.Dump()
}
func (b *StatsBucket) VBHash(docID string) uint32 {
	return b.bucket.VBHash(docID)
}
<commit_msg>Add sg-bucket interface methods to stats bucket<commit_after>package base

import (
	"encoding/json"
	"fmt"
	"sync/atomic"

	sgbucket "github.com/couchbase/sg-bucket"
)

// A wrapper around a Bucket that tracks bucket usage statistics as basic read/write counts.  Doesn't break
// down by operation type, to better identify counts for bulk operations
type StatsBucket struct {
	bucket             Bucket
	docsRead           uint64
	docsWritten        uint64
	bytesRead          uint64
	bytesWritten       uint64
	unknownSizeRead    uint64
	unknownSizeWritten uint64
}

type StatsBucketStats struct {
	DocsRead           uint64
	DocsWritten        uint64
	BytesRead          uint64
	BytesWritten       uint64
	UnknownSizeRead    uint64
	UnknownSizeWritten uint64
}

func (sbs *StatsBucketStats) String() string {
	return fmt.Sprintf("\nDocs Read:          %12d \nDocs Written:       %12d \nBytes Read:         %12d \nBytes Written:      %12d \nUnknown Size Reads: %12d \nUnknown Size Writes:%12d",
		sbs.DocsRead, sbs.DocsWritten, sbs.BytesRead, sbs.BytesWritten, sbs.UnknownSizeRead, sbs.UnknownSizeWritten)
}

func (sbs *StatsBucketStats) PerIteration(iterationCount uint64) string {
	return fmt.Sprintf("\nDocs Read:          %12d \nDocs Written:       %12d \nBytes Read:         %12d \nBytes Written:      %12d \nUnknown Size Reads: %12d \nUnknown Size Writes:%12d",
		sbs.DocsRead/iterationCount,
		sbs.DocsWritten/iterationCount,
		sbs.BytesRead/iterationCount,
		sbs.BytesWritten/iterationCount,
		sbs.UnknownSizeRead/iterationCount,
		sbs.UnknownSizeWritten/iterationCount)
}

func NewStatsBucket(bucket Bucket) *StatsBucket {
	return &StatsBucket{
		bucket: bucket,
	}
}

func (b *StatsBucket) docRead(count, bytesRead int) {
	atomic.AddUint64(&b.docsRead, uint64(count))
	if bytesRead == -1 {
		atomic.AddUint64(&b.unknownSizeRead, uint64(1))
	} else {
		atomic.AddUint64(&b.bytesRead, uint64(bytesRead))
	}
}

func (b *StatsBucket) docWrite(count, bytesWritten int) {
	atomic.AddUint64(&b.docsWritten, uint64(count))
	if bytesWritten == -1 {
		atomic.AddUint64(&b.unknownSizeWritten, uint64(1))
	} else {
		atomic.AddUint64(&b.bytesWritten, uint64(bytesWritten))

	}
}

func (b *StatsBucket) GetStats() StatsBucketStats {
	return StatsBucketStats{
		DocsRead:           atomic.LoadUint64(&b.docsRead),
		DocsWritten:        atomic.LoadUint64(&b.docsWritten),
		BytesRead:          atomic.LoadUint64(&b.bytesRead),
		BytesWritten:       atomic.LoadUint64(&b.bytesWritten),
		UnknownSizeRead:    atomic.LoadUint64(&b.unknownSizeRead),
		UnknownSizeWritten: atomic.LoadUint64(&b.unknownSizeWritten),
	}
}

func (b *StatsBucket) GetName() string {
	return b.bucket.GetName()
}
func (b *StatsBucket) Get(k string, rv interface{}) (uint64, error) {

	cas, err := b.bucket.Get(k, rv)
	if vBytes, ok := rv.([]byte); ok {
		defer b.docRead(1, len(vBytes))
	} else if marshalledJSON, marshalErr := json.Marshal(rv); marshalErr == nil {
		defer b.docRead(1, len(marshalledJSON))
	} else {
		defer b.docRead(1, -1)
	}
	return cas, err
}
func (b *StatsBucket) GetRaw(k string) (v []byte, cas uint64, err error) {
	v, cas, err = b.bucket.GetRaw(k)
	b.docRead(1, len(v))
	return v, cas, err
}
func (b *StatsBucket) GetAndTouchRaw(k string, exp int) (v []byte, cas uint64, err error) {
	v, cas, err = b.bucket.GetAndTouchRaw(k, exp)
	b.docRead(1, len(v))
	return v, cas, err
}
func (b *StatsBucket) GetBulkRaw(keys []string) (map[string][]byte, error) {
	results, err := b.bucket.GetBulkRaw(keys)
	for _, value := range results {
		b.docRead(1, len(value))
	}
	return results, err
}
func (b *StatsBucket) Add(k string, exp int, v interface{}) (added bool, err error) {
	if vBytes, ok := v.([]byte); ok {
		defer b.docWrite(1, len(vBytes))
	} else {
		defer b.docWrite(1, -1)
	}
	return b.bucket.Add(k, exp, v)
}
func (b *StatsBucket) AddRaw(k string, exp int, v []byte) (added bool, err error) {
	defer b.docWrite(1, len(v))
	return b.bucket.AddRaw(k, exp, v)
}
func (b *StatsBucket) Append(k string, data []byte) error {
	defer b.docWrite(1, len(data))
	return b.bucket.Append(k, data)
}
func (b *StatsBucket) Set(k string, exp int, v interface{}) error {
	if vBytes, ok := v.([]byte); ok {
		defer b.docWrite(1, len(vBytes))
	} else {
		defer b.docWrite(1, -1)
	}
	return b.bucket.Set(k, exp, v)
}
func (b *StatsBucket) SetRaw(k string, exp int, v []byte) error {
	defer b.docWrite(1, len(v))
	return b.bucket.SetRaw(k, exp, v)
}
func (b *StatsBucket) Delete(k string) error {
	return b.bucket.Delete(k)
}
func (b *StatsBucket) Write(k string, flags int, exp int, v interface{}, opt sgbucket.WriteOptions) error {
	if vBytes, ok := v.([]byte); ok {
		defer b.docWrite(1, len(vBytes))
	} else {
		defer b.docWrite(1, -1)
	}
	return b.bucket.Write(k, flags, exp, v, opt)
}
func (b *StatsBucket) WriteCas(k string, flags int, exp int, cas uint64, v interface{}, opt sgbucket.WriteOptions) (uint64, error) {
	if vBytes, ok := v.([]byte); ok {
		defer b.docWrite(1, len(vBytes))
	} else {
		defer b.docWrite(1, -1)
	}
	return b.bucket.WriteCas(k, flags, exp, cas, v, opt)
}
func (b *StatsBucket) Update(k string, exp int, callback sgbucket.UpdateFunc) (err error) {
	defer b.docWrite(1, -1)
	return b.bucket.Update(k, exp, callback)
}
func (b *StatsBucket) WriteUpdate(k string, exp int, callback sgbucket.WriteUpdateFunc) (err error) {
	defer b.docWrite(1, -1)
	return b.bucket.WriteUpdate(k, exp, callback)
}
func (b *StatsBucket) Incr(k string, amt, def uint64, exp int) (uint64, error) {
	return b.bucket.Incr(k, amt, def, exp)
}
func (b *StatsBucket) GetDDoc(docname string, value interface{}) error {
	return b.bucket.GetDDoc(docname, value)
}
func (b *StatsBucket) PutDDoc(docname string, value interface{}) error {
	return b.bucket.PutDDoc(docname, value)
}
func (b *StatsBucket) DeleteDDoc(docname string) error {
	return b.bucket.DeleteDDoc(docname)
}
func (b *StatsBucket) View(ddoc, name string, params map[string]interface{}) (sgbucket.ViewResult, error) {
	return b.bucket.View(ddoc, name, params)
}

func (b *StatsBucket) ViewCustom(ddoc, name string, params map[string]interface{}, vres interface{}) error {
	return b.bucket.ViewCustom(ddoc, name, params, vres)
}

func (b *StatsBucket) SetBulk(entries []*sgbucket.BulkSetEntry) (err error) {
	defer b.docWrite(len(entries), 0)
	return b.bucket.SetBulk(entries)
}

func (b *StatsBucket) Refresh() error {
	return b.bucket.Refresh()
}

func (b *StatsBucket) StartTapFeed(args sgbucket.TapArguments) (sgbucket.TapFeed, error) {
	return b.bucket.StartTapFeed(args)
}
func (b *StatsBucket) Close() {
	b.bucket.Close()
}
func (b *StatsBucket) Dump() {
	b.bucket.Dump()
}
func (b *StatsBucket) VBHash(docID string) uint32 {
	return b.bucket.VBHash(docID)
}

func (b *StatsBucket) GetMaxVbno() (uint16, error) {
	return b.bucket.GetMaxVbno()
}

func (b *StatsBucket) CouchbaseServerVersion() (major uint64, minor uint64, micro string, err error) {
	return b.bucket.CouchbaseServerVersion()
}

func (b *StatsBucket) UUID() (string, error) {
	return b.bucket.UUID()
}<|endoftext|>