<commit_before>package pushbullet

import (
	"bytes"
	"encoding/json"
	"errors"
	"net/http"
)

var (
	// API is the Pushbullet API endpoint.
	API = "https://api.pushbullet.com/v2/pushes"
)

type apiResponse struct {
	Active                  bool    `json:"active"`
	Iden                    string  `json:"iden"`
	Created                 float64 `json:"created"`
	Modified                float64 `json:"modified"`
	Type                    string  `json:"type"`
	Dismissed               bool    `json:"dismissed"`
	Direction               string  `json:"direction"`
	SenderIden              string  `json:"sender_iden"`
	SenderEmail             string  `json:"sender_email"`
	SenderEmailNormalized   string  `json:"sender_email_normalized"`
	SenderName              string  `json:"sender_name"`
	ReceiverIden            string  `json:"receiver_iden"`
	ReceiverEmail           string  `json:"receiver_email"`
	ReceiverEmailNormalized string  `json:"receiver_email_normalized"`
	Title                   string  `json:"title"`
	Body                    string  `json:"body"`
	Error                   struct {
		Code    string `json:"code"`
		Type    string `json:"type"`
		Message string `json:"message"`
		Cat     string `json:"cat"`
	} `json:"error"`
	ErrorCode string `json:"error_code"`
}

// Notification is a pushbullet notification.
type Notification struct {
	Body        string `json:"body"`
	Title       string `json:"title"`
	Type        string `json:"type"`
	DeviceIden  string `json:"device_iden"`

	AccessToken string       `json:"-"`
	Client      *http.Client `json:"-"`
}

// Send sends a Pushbullet notification.
func (n *Notification) Send() error {
	payload := new(bytes.Buffer)
	if err := json.NewEncoder(payload).Encode(n); err != nil {
		return err
	}

	req, err := http.NewRequest("POST", API, payload)
	if err != nil {
		return err
	}
	req.Header.Set("Access-Token", n.AccessToken)
	req.Header.Set("Content-Type", "application/json")

	resp, err := n.Client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	var r apiResponse
	if err := json.NewDecoder(resp.Body).Decode(&r); err != nil {
		return err
	}

	if r.ErrorCode != "" {
		return errors.New(r.ErrorCode)
	}

	return nil
}
<commit_msg>Run gofmt on project (#81)<commit_after>package pushbullet

import (
	"bytes"
	"encoding/json"
	"errors"
	"net/http"
)

var (
	// API is the Pushbullet API endpoint.
	API = "https://api.pushbullet.com/v2/pushes"
)

type apiResponse struct {
	Active                  bool    `json:"active"`
	Iden                    string  `json:"iden"`
	Created                 float64 `json:"created"`
	Modified                float64 `json:"modified"`
	Type                    string  `json:"type"`
	Dismissed               bool    `json:"dismissed"`
	Direction               string  `json:"direction"`
	SenderIden              string  `json:"sender_iden"`
	SenderEmail             string  `json:"sender_email"`
	SenderEmailNormalized   string  `json:"sender_email_normalized"`
	SenderName              string  `json:"sender_name"`
	ReceiverIden            string  `json:"receiver_iden"`
	ReceiverEmail           string  `json:"receiver_email"`
	ReceiverEmailNormalized string  `json:"receiver_email_normalized"`
	Title                   string  `json:"title"`
	Body                    string  `json:"body"`
	Error                   struct {
		Code    string `json:"code"`
		Type    string `json:"type"`
		Message string `json:"message"`
		Cat     string `json:"cat"`
	} `json:"error"`
	ErrorCode string `json:"error_code"`
}

// Notification is a pushbullet notification.
type Notification struct {
	Body       string `json:"body"`
	Title      string `json:"title"`
	Type       string `json:"type"`
	DeviceIden string `json:"device_iden"`

	AccessToken string       `json:"-"`
	Client      *http.Client `json:"-"`
}

// Send sends a Pushbullet notification.
func (n *Notification) Send() error {
	payload := new(bytes.Buffer)
	if err := json.NewEncoder(payload).Encode(n); err != nil {
		return err
	}

	req, err := http.NewRequest("POST", API, payload)
	if err != nil {
		return err
	}
	req.Header.Set("Access-Token", n.AccessToken)
	req.Header.Set("Content-Type", "application/json")

	resp, err := n.Client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	var r apiResponse
	if err := json.NewDecoder(resp.Body).Decode(&r); err != nil {
		return err
	}

	if r.ErrorCode != "" {
		return errors.New(r.ErrorCode)
	}

	return nil
}
<|endoftext|>


<commit_before>// AnsiGo 1.00 (c) by Frederic Cambus 2012
// http://www.github.com/fcambus/ansigo
//
// Created:      2012/02/14
// Last Updated: 2012/02/19
//
// AnsiGo is released under the MIT license.
// See LICENSE file for details.
package main

import (
	"fmt"
	"image"
	"image/draw"
	"image/png"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
)

func main() {

	fmt.Println("-------------------------------------------------------------------------------\n                    AnsiGo 1.00 (c) by Frederic CAMBUS 2012\n-------------------------------------------------------------------------------\n")

	// Check input parameters and show usage
	if len(os.Args) != 2 {
		fmt.Println("USAGE:    ansigo inputfile\n")
		fmt.Println("EXAMPLES: ansigo ansi.ans\n")
		os.Exit(1)
	}

	input := os.Args[1]
	output := input + ".png"

	fmt.Println("Input File:", input)
	fmt.Println("Output File:", output)

	var ansi Ansi
	ansi.SetPalette()
	ansi.SetFont()

	// Load input file
	data, err := ioutil.ReadFile(input)
	if err != nil {
		fmt.Println("\nERROR: Can't open or read", input, "\n")
		os.Exit(1)
	}

	// Process ANSI
	for i := 0; i < len(data); i++ {
		ansi.character = data[i]

		// 80th column wrapping
		if ansi.positionX == 80 {
			ansi.positionY++
			ansi.positionX = 0
		}

		// CR (Carriage Return)
		if ansi.character == '\r' {
			continue
		}

		// LF (Line Feed)
		if ansi.character == '\n' {
			ansi.positionY++
			ansi.positionX = 0
			continue
		}

		// HT (Horizontal Tabulation)
		if ansi.character == '\t' {
			ansi.positionX += 8
			continue
		}

		// SUB (Substitute)
		if ansi.character == '\x1a' {
			break
		}

		// ANSI Sequence : ESC (Escape) + [
		if ansi.character == '\x1b' && data[i+1] == '[' {
			ansiSequence := []byte{}

			for j := 0; j < 12; j++ {
				ansiSequenceCharacter := data[i+2+j]

				// Cursor Position
				if ansiSequenceCharacter == 'H' || ansiSequenceCharacter == 'f' {
					ansiSequenceValues := strings.SplitN(string(ansiSequence), ";", -1)

					valueY, _ := strconv.Atoi(ansiSequenceValues[0])
					ansi.positionY = valueY - 1

					valueX, _ := strconv.Atoi(ansiSequenceValues[1])
					ansi.positionX = valueX - 1

					i += j + 2
					break
				}

				// Cursor Up
				if ansiSequenceCharacter == 'A' {
					valueY, _ := strconv.Atoi(string(ansiSequence))
					if valueY == 0 {
						valueY++
					}

					ansi.positionY = ansi.positionY - valueY

					i += j + 2
					break
				}

				// Cursor Down
				if ansiSequenceCharacter == 'B' {
					valueY, _ := strconv.Atoi(string(ansiSequence))
					if valueY == 0 {
						valueY++
					}

					ansi.positionY = ansi.positionY + valueY

					i += j + 2
					break
				}

				// Cursor Forward
				if ansiSequenceCharacter == 'C' {
					valueX, _ := strconv.Atoi(string(ansiSequence))
					if valueX == 0 {
						valueX++
					}

					ansi.positionX = ansi.positionX + valueX
					if ansi.positionX > 80 {
						ansi.positionX = 80
					}

					i += j + 2
					break
				}

				// Cursor Backward
				if ansiSequenceCharacter == 'D' {
					valueX, _ := strconv.Atoi(string(ansiSequence))
					if valueX == 0 {
						valueX++
					}

					ansi.positionX = ansi.positionX - valueX
					if ansi.positionX < 0 {
						ansi.positionX = 0
					}

					i += j + 2
					break
				}

				// Save Cursor Position
				if ansiSequenceCharacter == 's' {
					ansi.savedPositionY = ansi.positionY
					ansi.savedPositionX = ansi.positionX

					i += j + 2
					break
				}

				// Restore Cursor Position
				if ansiSequenceCharacter == 'u' {
					ansi.positionY = ansi.savedPositionY
					ansi.positionX = ansi.savedPositionX

					i += j + 2
					break
				}

				// Erase Display
				if ansiSequenceCharacter == 'J' {
					value, _ := strconv.Atoi(string(ansiSequence))

					if value == 2 {
						ansi.buffer = nil

						ansi.positionX = 0
						ansi.positionY = 0
						ansi.sizeX = 0
						ansi.sizeY = 0
					}

					i += j + 2
					break
				}

				// Set Graphic Rendition
				if ansiSequenceCharacter == 'm' {
					ansiSequenceValues := strings.SplitN(string(ansiSequence), ";", -1)

					for j := 0; j < len(ansiSequenceValues); j++ {
						valueColor, _ := strconv.Atoi(ansiSequenceValues[j])

						switch valueColor {
						case 0:
							ansi.colorBackground = 0
							ansi.colorForeground = 7
							ansi.bold = false

						case 1:
							ansi.colorForeground += 8
							ansi.bold = true

						case 5:
							ansi.colorBackground += 8

						case 30, 31, 32, 33, 34, 35, 36, 37:
							ansi.colorForeground = valueColor - 30
							if ansi.bold {
								ansi.colorForeground += 8
							}

						case 40, 41, 42, 43, 44, 45, 46, 47:
							ansi.colorBackground = valueColor - 40
						}
					}

					i += j + 2
					break
				}

				// Skipping Set Mode And Reset Mode Sequences
				if ansiSequenceCharacter == 'h' || ansiSequenceCharacter == 'l' {

					i += j + 2
					break
				}

				ansiSequence = append(ansiSequence, ansiSequenceCharacter)
			}
		} else {
			// Record Number Of Columns And Lines Used
			if ansi.positionX > ansi.sizeX {
				ansi.sizeX = ansi.positionX
			}
			if ansi.positionY > ansi.sizeY {
				ansi.sizeY = ansi.positionY
			}

			// Write Current Character Info In A Temporary Array
			ansi.buffer = append(ansi.buffer, Character{ansi.colorBackground, ansi.colorForeground, ansi.positionX, ansi.positionY, ansi.character})

			ansi.positionX++
		}
	}

	// Allocate Image Buffer Memory
	canvasSize := image.Rect(0, 0, 640, (ansi.sizeY+1)*16)
	canvas := image.NewRGBA(canvasSize)

	// Draw The Canvas Background
	draw.Draw(canvas, canvas.Bounds(), &image.Uniform{ansi.palette[0]}, image.ZP, draw.Src)

	// Render ANSI
	for i := 0; i < len(ansi.buffer); i++ {
		character := ansi.buffer[i]

		// Set Background
		draw.Draw(canvas, image.Rect(character.positionX*8, character.positionY*16, character.positionX*8+8, character.positionY*16+16), &image.Uniform{ansi.palette[character.colorBackground]}, image.ZP, draw.Src)

		// Draw Character
		for line := 0; line < 16; line++ {
			for column := 0; column < 8; column++ {
				if (ansi.font[line+(int(character.code)*16)] & (0x80 >> uint(column))) != 0 {
					canvas.Set(character.positionX*8+column, line+character.positionY*16, ansi.palette[character.colorForeground])
				}
			}
		}
	}

	// Create Output File
	outputFile, err := os.Create(output)
	if err != nil {
		fmt.Println("ERROR: Can't create ouput file.")
		os.Exit(1)
	}

	// Encode PNG image
	if err = png.Encode(outputFile, canvas); err != nil {
		fmt.Println("ERROR: Can't encode PNG file.")
		os.Exit(1)
	}

	outputFile.Close()

	fmt.Println("\nSuccessfully created file", output, "\n")
}
<commit_msg>Modifying ANSI sequence parser to use switch<commit_after>// AnsiGo 1.00 (c) by Frederic Cambus 2012
// http://www.github.com/fcambus/ansigo
//
// Created:      2012/02/14
// Last Updated: 2012/02/19
//
// AnsiGo is released under the MIT license.
// See LICENSE file for details.
package main

import (
	"fmt"
	"image"
	"image/draw"
	"image/png"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
)

func main() {

	fmt.Println("-------------------------------------------------------------------------------\n                    AnsiGo 1.00 (c) by Frederic CAMBUS 2012\n-------------------------------------------------------------------------------\n")

	// Check input parameters and show usage
	if len(os.Args) != 2 {
		fmt.Println("USAGE:    ansigo inputfile\n")
		fmt.Println("EXAMPLES: ansigo ansi.ans\n")
		os.Exit(1)
	}

	input := os.Args[1]
	output := input + ".png"

	fmt.Println("Input File:", input)
	fmt.Println("Output File:", output)

	var ansi Ansi
	ansi.SetPalette()
	ansi.SetFont()

	// Load input file
	data, err := ioutil.ReadFile(input)
	if err != nil {
		fmt.Println("\nERROR: Can't open or read", input, "\n")
		os.Exit(1)
	}

	// Process ANSI
	for i := 0; i < len(data); i++ {
		ansi.character = data[i]

		// 80th column wrapping
		if ansi.positionX == 80 {
			ansi.positionY++
			ansi.positionX = 0
		}

		// CR (Carriage Return)
		if ansi.character == '\r' {
			continue
		}

		// LF (Line Feed)
		if ansi.character == '\n' {
			ansi.positionY++
			ansi.positionX = 0
			continue
		}

		// HT (Horizontal Tabulation)
		if ansi.character == '\t' {
			ansi.positionX += 8
			continue
		}

		// SUB (Substitute)
		if ansi.character == '\x1a' {
			break
		}

		// ANSI Sequence : ESC (Escape) + [
		if ansi.character == '\x1b' && data[i+1] == '[' {
			ansiSequence := []byte{}

		sequence:
			for j := 0; j < 12; j++ {
				ansiSequenceCharacter := data[i+2+j]

				switch ansiSequenceCharacter {
				case 'H', 'f': // Cursor Position
					ansiSequenceValues := strings.SplitN(string(ansiSequence), ";", -1)

					valueY, _ := strconv.Atoi(ansiSequenceValues[0])
					ansi.positionY = valueY - 1

					valueX, _ := strconv.Atoi(ansiSequenceValues[1])
					ansi.positionX = valueX - 1

					i += j + 2
					break sequence

				case 'A': // Cursor Up
					valueY, _ := strconv.Atoi(string(ansiSequence))
					if valueY == 0 {
						valueY++
					}

					ansi.positionY = ansi.positionY - valueY

					i += j + 2
					break sequence

				case 'B': // Cursor Down
					valueY, _ := strconv.Atoi(string(ansiSequence))
					if valueY == 0 {
						valueY++
					}

					ansi.positionY = ansi.positionY + valueY

					i += j + 2
					break sequence

				case 'C': // Cursor Forward
					valueX, _ := strconv.Atoi(string(ansiSequence))
					if valueX == 0 {
						valueX++
					}

					ansi.positionX = ansi.positionX + valueX
					if ansi.positionX > 80 {
						ansi.positionX = 80
					}

					i += j + 2
					break sequence

				case 'D': // Cursor Backward
					valueX, _ := strconv.Atoi(string(ansiSequence))
					if valueX == 0 {
						valueX++
					}

					ansi.positionX = ansi.positionX - valueX
					if ansi.positionX < 0 {
						ansi.positionX = 0
					}

					i += j + 2
					break sequence

				case 's': // Save Cursor Position
					ansi.savedPositionY = ansi.positionY
					ansi.savedPositionX = ansi.positionX

					i += j + 2
					break sequence

				case 'u': // Restore Cursor Position
					ansi.positionY = ansi.savedPositionY
					ansi.positionX = ansi.savedPositionX

					i += j + 2
					break sequence

				case 'J': // Erase Display
					value, _ := strconv.Atoi(string(ansiSequence))

					if value == 2 {
						ansi.buffer = nil

						ansi.positionX = 0
						ansi.positionY = 0
						ansi.sizeX = 0
						ansi.sizeY = 0
					}

					i += j + 2
					break sequence

				case 'm': // Set Graphic Rendition
					ansiSequenceValues := strings.SplitN(string(ansiSequence), ";", -1)

					for j := 0; j < len(ansiSequenceValues); j++ {
						valueColor, _ := strconv.Atoi(ansiSequenceValues[j])

						switch valueColor {
						case 0:
							ansi.colorBackground = 0
							ansi.colorForeground = 7
							ansi.bold = false

						case 1:
							ansi.colorForeground += 8
							ansi.bold = true

						case 5:
							ansi.colorBackground += 8

						case 30, 31, 32, 33, 34, 35, 36, 37:
							ansi.colorForeground = valueColor - 30
							if ansi.bold {
								ansi.colorForeground += 8
							}

						case 40, 41, 42, 43, 44, 45, 46, 47:
							ansi.colorBackground = valueColor - 40
						}
					}

					i += j + 2
					break sequence

				case 'h', 'l': // Skipping Set Mode And Reset Mode Sequences
					i += j + 2
					break sequence
				}

				ansiSequence = append(ansiSequence, ansiSequenceCharacter)
			}
		} else {
			// Record Number Of Columns And Lines Used
			if ansi.positionX > ansi.sizeX {
				ansi.sizeX = ansi.positionX
			}
			if ansi.positionY > ansi.sizeY {
				ansi.sizeY = ansi.positionY
			}

			// Write Current Character Info In A Temporary Array
			ansi.buffer = append(ansi.buffer, Character{ansi.colorBackground, ansi.colorForeground, ansi.positionX, ansi.positionY, ansi.character})

			ansi.positionX++
		}
	}

	// Allocate Image Buffer Memory
	canvasSize := image.Rect(0, 0, 640, (ansi.sizeY+1)*16)
	canvas := image.NewRGBA(canvasSize)

	// Draw The Canvas Background
	draw.Draw(canvas, canvas.Bounds(), &image.Uniform{ansi.palette[0]}, image.ZP, draw.Src)

	// Render ANSI
	for i := 0; i < len(ansi.buffer); i++ {
		character := ansi.buffer[i]

		// Set Background
		draw.Draw(canvas, image.Rect(character.positionX*8, character.positionY*16, character.positionX*8+8, character.positionY*16+16), &image.Uniform{ansi.palette[character.colorBackground]}, image.ZP, draw.Src)

		// Draw Character
		for line := 0; line < 16; line++ {
			for column := 0; column < 8; column++ {
				if (ansi.font[line+(int(character.code)*16)] & (0x80 >> uint(column))) != 0 {
					canvas.Set(character.positionX*8+column, line+character.positionY*16, ansi.palette[character.colorForeground])
				}
			}
		}
	}

	// Create Output File
	outputFile, err := os.Create(output)
	if err != nil {
		fmt.Println("ERROR: Can't create ouput file.")
		os.Exit(1)
	}

	// Encode PNG image
	if err = png.Encode(outputFile, canvas); err != nil {
		fmt.Println("ERROR: Can't encode PNG file.")
		os.Exit(1)
	}

	outputFile.Close()

	fmt.Println("\nSuccessfully created file", output, "\n")
}
<|endoftext|>


<commit_before>// +build linux,unit

// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 metrics

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/config"
	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	"github.com/stretchr/testify/assert"
)

// Create default config for Metrics. PrometheusMetricsEnabled is set to false
// by default, so it must be turned on
func getTestConfig() config.Config {
	cfg := config.DefaultConfig()
	cfg.PrometheusMetricsEnabled = true
	return cfg
}

// Tests if MetricsEngineGlobal variable is initialized and if all managed
// MetricsClients are initialized
func TestMetricsEngineInit(t *testing.T) {
	defer func() {
		MetricsEngineGlobal = &MetricsEngine{
			collection: false,
		}
	}()
	cfg := getTestConfig()
	MustInit(&cfg, prometheus.NewRegistry())
	assert.NotNil(t, MetricsEngineGlobal)
	assert.Equal(t, len(MetricsEngineGlobal.managedMetrics), len(managedAPIs))
}

// Tests if a default config will start Prometheus metrics. Should be disabled
// by default.
func TestDisablePrometheusMetrics(t *testing.T) {
	defer func() {
		MetricsEngineGlobal = &MetricsEngine{
			collection: false,
		}
	}()
	cfg := getTestConfig()
	cfg.PrometheusMetricsEnabled = false
	MustInit(&cfg, prometheus.NewRegistry())
	PublishMetrics()
	assert.False(t, MetricsEngineGlobal.collection)
}

// Mimicks metric collection of Docker API calls through Go routines. The method
// call to record a metric is the same used by various clients throughout Agent.
// We sleep the go routine to simulate "work" being done.
// We can determine the expected values for the metrics and create a map of them,
// which will then be used to verify the accuracy of the metrics collected.
func TestMetricCollection(t *testing.T) {
	defer func() {
		MetricsEngineGlobal = &MetricsEngine{
			collection: false,
		}
	}()
	cfg := getTestConfig()
	MustInit(&cfg, prometheus.NewRegistry())
	MetricsEngineGlobal.collection = true

	var DockerMetricSleepTime4 time.Duration = 4 * time.Second
	var DockerMetricSleepTime2 time.Duration = 2 * time.Second

	var wg sync.WaitGroup
	wg.Add(40)

	// These Go routines simulate metrics collection
	go func() {
		for i := 0; i < 10; i++ {
			go func() {
				defer wg.Done()
				defer MetricsEngineGlobal.RecordDockerMetric("START")()
				time.Sleep(DockerMetricSleepTime4)
			}()
		}
	}()
	go func() {
		for i := 0; i < 10; i++ {
			go func() {
				defer wg.Done()
				defer MetricsEngineGlobal.RecordDockerMetric("START")()
				time.Sleep(DockerMetricSleepTime2)
			}()
		}
	}()
	go func() {
		for i := 0; i < 10; i++ {
			go func() {
				defer wg.Done()
				defer MetricsEngineGlobal.RecordDockerMetric("STOP")()
				time.Sleep(DockerMetricSleepTime4)
			}()
		}
	}()
	go func() {
		for i := 0; i < 10; i++ {
			go func() {
				defer wg.Done()
				defer MetricsEngineGlobal.RecordDockerMetric("STOP")()
				time.Sleep(DockerMetricSleepTime2)
			}()
		}
	}()
	time.Sleep(2 * time.Second)
	wg.Wait()

	// This will gather all collected metrics and store them in a MetricFamily list
	// All metric families can be printed by looping over this variable using
	// fmt.Println(proto.MarshalTextString(metricFamilies[n])) where n = index
	metricFamilies, err := MetricsEngineGlobal.Registry.Gather()
	assert.NoError(t, err)

	// Here we set up the expectations. These are known values which make verfication
	// easier.
	expected := make(metricMap)
	expected["AgentMetrics_DockerAPI_call_count"] = make(map[string][]interface{})
	expected["AgentMetrics_DockerAPI_call_count"]["CallSTART"] = []interface{}{
		"COUNTER",
		20.0,
	}
	expected["AgentMetrics_DockerAPI_call_count"]["CallSTOP"] = []interface{}{
		"COUNTER",
		20.0,
	}
	expected["AgentMetrics_DockerAPI_call_duration"] = make(map[string][]interface{})
	expected["AgentMetrics_DockerAPI_call_duration"]["CallSTART"] = []interface{}{
		"GUAGE",
		0.0,
	}
	expected["AgentMetrics_DockerAPI_call_duration"]["CallSTOP"] = []interface{}{
		"GUAGE",
		0.0,
	}
	expected["AgentMetrics_DockerAPI_duration_seconds"] = make(map[string][]interface{})
	expected["AgentMetrics_DockerAPI_duration_seconds"]["CallSTART"] = []interface{}{
		"SUMMARY",
		3.0,
	}
	expected["AgentMetrics_DockerAPI_duration_seconds"]["CallSTOP"] = []interface{}{
		"SUMMARY",
		3.0,
	}
	// We will do a simple tree search to verify all metrics in metricsFamilies
	// are as expected
	assert.True(t, verifyStats(metricFamilies, expected), "Metrics are not accurate")
}

// A type for storing a Tree-based map. We map the MetricName to a map of metrics
// under that name. This second map indexes by MetricLabelName+MetricLabelValue to
// a slice MetricType and MetricValue.
//MetricName:metricLabelName+metricLabelValue:[metricType, metricValue]
type metricMap map[string]map[string][]interface{}

// In order to verify the MetricFamily with the expected metric values, we do a simple
// tree search to verify that all stats in the MetricFamily coincide with the expected
// metric values.
// This method only verifes that all metrics in var metricsReceived are present in
// var expectedMetrics
func verifyStats(metricsReceived []*dto.MetricFamily, expectedMetrics metricMap) bool {
	var threshhold float64 = 0.1 // Maximum threshhold for two metrics being equal
	for _, metricFamily := range metricsReceived {
		if metricList, found := expectedMetrics[metricFamily.GetName()]; found {
			for _, metric := range metricFamily.GetMetric() {
				if aMetric, found := metricList[metric.GetLabel()[0].GetName()+metric.GetLabel()[0].GetValue()]; found {
					metricTypeExpected := string(aMetric[0].(string))
					metricValExpected := float64(aMetric[1].(float64))
					switch metricTypeExpected {
					case "GUAGE":
						continue
					case "COUNTER":
						if !compareDiff(metricValExpected, metric.GetCounter().GetValue(), threshhold) {
							fmt.Printf("Does not match SUMMARY. Expected: %f, Received: %f\n", metricValExpected, metric.GetCounter().GetValue())
							return false
						}
					case "SUMMARY":
						if !compareDiff(metricValExpected, metric.GetSummary().GetSampleSum()/float64(metric.GetSummary().GetSampleCount()), threshhold) {
							fmt.Printf("Does not match SUMMARY. Expected: %f, Received: %f\n", metricValExpected, metric.GetSummary().GetSampleSum()/float64(metric.GetSummary().GetSampleCount()))
							return false
						}
					default:
						fmt.Println("Metric Type not recognized")
						return false
					}
				} else {
					fmt.Println("MetricLabel Name and Value combo not found")
					return false
				}
			}
		} else {
			fmt.Println("MetricName not found")
			return false
		}
	}
	return true
}

// Helper function to determine if two values (a and b) are within a percentage of a
func compareDiff(a, b, deltaMin float64) bool {
	diff := a - b
	if diff < 0 {
		diff = -diff
	}
	return diff <= (a * deltaMin)
}
<commit_msg>Flaky unit test: metrics collection (#2223)<commit_after>// +build linux,unit

// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 metrics

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/config"
	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
	"github.com/stretchr/testify/assert"
)

// Create default config for Metrics. PrometheusMetricsEnabled is set to false
// by default, so it must be turned on
func getTestConfig() config.Config {
	cfg := config.DefaultConfig()
	cfg.PrometheusMetricsEnabled = true
	return cfg
}

// Tests if MetricsEngineGlobal variable is initialized and if all managed
// MetricsClients are initialized
func TestMetricsEngineInit(t *testing.T) {
	defer func() {
		MetricsEngineGlobal = &MetricsEngine{
			collection: false,
		}
	}()
	cfg := getTestConfig()
	MustInit(&cfg, prometheus.NewRegistry())
	assert.NotNil(t, MetricsEngineGlobal)
	assert.Equal(t, len(MetricsEngineGlobal.managedMetrics), len(managedAPIs))
}

// Tests if a default config will start Prometheus metrics. Should be disabled
// by default.
func TestDisablePrometheusMetrics(t *testing.T) {
	defer func() {
		MetricsEngineGlobal = &MetricsEngine{
			collection: false,
		}
	}()
	cfg := getTestConfig()
	cfg.PrometheusMetricsEnabled = false
	MustInit(&cfg, prometheus.NewRegistry())
	PublishMetrics()
	assert.False(t, MetricsEngineGlobal.collection)
}

// Mimicks metric collection of Docker API calls through Go routines. The method
// call to record a metric is the same used by various clients throughout Agent.
// We sleep the go routine to simulate "work" being done.
// We can determine the expected values for the metrics and create a map of them,
// which will then be used to verify the accuracy of the metrics collected.
func TestMetricCollection(t *testing.T) {
	defer func() {
		MetricsEngineGlobal = &MetricsEngine{
			collection: false,
		}
	}()
	cfg := getTestConfig()
	MustInit(&cfg, prometheus.NewRegistry())
	MetricsEngineGlobal.collection = true

	var DockerMetricSleepTime1 time.Duration = 1 * time.Second
	var DockerMetricSleepTime2 time.Duration = 2 * time.Second

	var wg sync.WaitGroup
	wg.Add(20)

	// These Go routines simulate metrics collection
	for i := 0; i < 5; i++ {
		go func() {
			defer wg.Done()
			defer MetricsEngineGlobal.RecordDockerMetric("START")()
			time.Sleep(DockerMetricSleepTime1)
		}()
	}
	for i := 0; i < 5; i++ {
		go func() {
			defer wg.Done()
			defer MetricsEngineGlobal.RecordDockerMetric("START")()
			time.Sleep(DockerMetricSleepTime2)
		}()
	}
	for i := 0; i < 5; i++ {
		go func() {
			defer wg.Done()
			defer MetricsEngineGlobal.RecordDockerMetric("STOP")()
			time.Sleep(DockerMetricSleepTime1)
		}()
	}
	for i := 0; i < 5; i++ {
		go func() {
			defer wg.Done()
			defer MetricsEngineGlobal.RecordDockerMetric("STOP")()
			time.Sleep(DockerMetricSleepTime2)
		}()
	}
	wg.Wait()
	time.Sleep(time.Second)

	// This will gather all collected metrics and store them in a MetricFamily list
	// All metric families can be printed by looping over this variable using
	// fmt.Println(proto.MarshalTextString(metricFamilies[n])) where n = index
	metricFamilies, err := MetricsEngineGlobal.Registry.Gather()
	assert.NoError(t, err)

	// Here we set up the expectations. These are known values which make verfication
	// easier.
	expected := make(metricMap)
	expected["AgentMetrics_DockerAPI_call_count"] = make(map[string][]interface{})
	expected["AgentMetrics_DockerAPI_call_count"]["CallSTART"] = []interface{}{
		"COUNTER",
		10.0,
	}
	expected["AgentMetrics_DockerAPI_call_count"]["CallSTOP"] = []interface{}{
		"COUNTER",
		10.0,
	}
	expected["AgentMetrics_DockerAPI_call_duration"] = make(map[string][]interface{})
	expected["AgentMetrics_DockerAPI_call_duration"]["CallSTART"] = []interface{}{
		"GUAGE",
		0.0,
	}
	expected["AgentMetrics_DockerAPI_call_duration"]["CallSTOP"] = []interface{}{
		"GUAGE",
		0.0,
	}
	expected["AgentMetrics_DockerAPI_duration_seconds"] = make(map[string][]interface{})
	expected["AgentMetrics_DockerAPI_duration_seconds"]["CallSTART"] = []interface{}{
		"SUMMARY",
		1.5,
	}
	expected["AgentMetrics_DockerAPI_duration_seconds"]["CallSTOP"] = []interface{}{
		"SUMMARY",
		1.5,
	}
	// We will do a simple tree search to verify all metrics in metricsFamilies
	// are as expected
	assert.True(t, verifyStats(metricFamilies, expected), "Metrics are not accurate")
}

// A type for storing a Tree-based map. We map the MetricName to a map of metrics
// under that name. This second map indexes by MetricLabelName+MetricLabelValue to
// a slice MetricType and MetricValue.
//MetricName:metricLabelName+metricLabelValue:[metricType, metricValue]
type metricMap map[string]map[string][]interface{}

// In order to verify the MetricFamily with the expected metric values, we do a simple
// tree search to verify that all stats in the MetricFamily coincide with the expected
// metric values.
// This method only verifes that all metrics in var metricsReceived are present in
// var expectedMetrics
func verifyStats(metricsReceived []*dto.MetricFamily, expectedMetrics metricMap) bool {
	var threshhold float64 = 0.1 // Maximum threshhold for two metrics being equal
	for _, metricFamily := range metricsReceived {
		if metricList, found := expectedMetrics[metricFamily.GetName()]; found {
			for _, metric := range metricFamily.GetMetric() {
				if aMetric, found := metricList[metric.GetLabel()[0].GetName()+metric.GetLabel()[0].GetValue()]; found {
					metricTypeExpected := string(aMetric[0].(string))
					metricValExpected := float64(aMetric[1].(float64))
					switch metricTypeExpected {
					case "GUAGE":
						continue
					case "COUNTER":
						if !compareDiff(metricValExpected, metric.GetCounter().GetValue(), threshhold) {
							fmt.Printf("Does not match SUMMARY. Expected: %f, Received: %f\n", metricValExpected, metric.GetCounter().GetValue())
							return false
						}
					case "SUMMARY":
						if !compareDiff(metricValExpected, metric.GetSummary().GetSampleSum()/float64(metric.GetSummary().GetSampleCount()), threshhold) {
							fmt.Printf("Does not match SUMMARY. Expected: %f, Received: %f\n", metricValExpected, metric.GetSummary().GetSampleSum()/float64(metric.GetSummary().GetSampleCount()))
							return false
						}
					default:
						fmt.Println("Metric Type not recognized")
						return false
					}
				} else {
					fmt.Println("MetricLabel Name and Value combo not found")
					return false
				}
			}
		} else {
			fmt.Println("MetricName not found")
			return false
		}
	}
	return true
}

// Helper function to determine if two values (a and b) are within a percentage of a
func compareDiff(a, b, deltaMin float64) bool {
	diff := a - b
	if diff < 0 {
		diff = -diff
	}
	return diff <= (a * deltaMin)
}
<|endoftext|>


<commit_before>package accumulate

import (
	"fmt"
	"strings"
	"testing"
)

func echo(c string) string {
	return c
}

var tests = []struct {
	expected    []string
	given       []string
	converter   func(string) string
	description string
}{
	{[]string{}, []string{}, echo, "echo"},
	{[]string{"HELLO", "WORLD"}, []string{"hello", "world"}, strings.ToUpper, "upcase"},
}

func TestAccumulate(t *testing.T) {
	for _, test := range tests {
		actual := Accumulate(test.given, test.converter)
		if fmt.Sprintf("%s", actual) != fmt.Sprintf("%s", test.expected) {
			t.Fatalf("Allergies(%s, %#v): expected %s, actual %s", test.given, test.converter, test.expected, actual)
		} else {
			t.Logf("PASS: %s %v", test.description, test.given)
		}
	}
}
<commit_msg>go/accumulate: changed 'Allergies' to 'Accumulate', added tests<commit_after>package accumulate

import (
	"fmt"
	"strings"
	"testing"
)

func echo(c string) string {
	return c
}

func capitalize(word string) string {
	return strings.Title(word)
}

var tests = []struct {
	expected    []string
	given       []string
	converter   func(string) string
	description string
}{
	{[]string{}, []string{}, echo, "echo"},
	{[]string{"echo", "echo", "echo", "echo"}, []string{"echo", "echo", "echo", "echo"}, echo, "echo"},
	{[]string{"First", "Letter", "Only"}, []string{"first", "letter", "only"}, capitalize, "capitalize"},
	{[]string{"HELLO", "WORLD"}, []string{"hello", "world"}, strings.ToUpper, "upcase"},
}

func TestAccumulate(t *testing.T) {
	for _, test := range tests {
		actual := Accumulate(test.given, test.converter)
		if fmt.Sprintf("%s", actual) != fmt.Sprintf("%s", test.expected) {
			t.Fatalf("Accumulate(%s, %#v): expected %s, actual %s", test.given, test.converter, test.expected, actual)
		} else {
			t.Logf("PASS: %s %v", test.description, test.given)
		}
	}
}
<|endoftext|>


<commit_before>package docli

import (
	"fmt"
	"strconv"

	"github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/digitalocean/godo"
)

// Get returns a droplet action by id.
func DropletActionGet(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	dropletID := c.Int(ArgDropletID)
	actionID := c.Int(ArgActionID)

	a, _, err := client.DropletActions.Get(dropletID, actionID)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not disable backups for droplet")
	}

	err = WriteJSON(a, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// DisableBackups disables backups for a droplet.
func DropletActionDisableBackups(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.DisableBackups(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not disable backups for droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Reboot reboots a droplet.
func DropletActionReboot(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.Reboot(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not reboot droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// PowerCycle power cycles a droplet.
func DropletActionPowerCycle(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	r, _, err := client.DropletActions.PowerCycle(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not power cycle droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Shutdown shuts a droplet down.
func DropletActionShutdown(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.Shutdown(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not shutdown droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// PowerOff turns droplet power off.
func DropletActionPowerOff(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.PowerOff(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not power off droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// PowerOn turns droplet power on.
func DropletActionPowerOn(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.PowerOn(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not power on droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// PasswordReset resets the droplet root password.
func DropletActionPasswordReset(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.PasswordReset(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not reset password for droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// EnableIPv6 enables IPv6 for a droplet.
func DropletActionEnableIPv6(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.EnableIPv6(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not enable IPv6 for droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// EnablePrivateNetworking enables private networking for a droplet.
func DropletActionEnablePrivateNetworking(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.EnablePrivateNetworking(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not enable private networking for droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Upgrade upgrades a droplet.
func DropletActionUpgrade(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.Upgrade(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not upgrade droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Restore restores a droplet using an image id.
func DropletActionRestore(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	image := c.Int(ArgImageID)

	r, _, err := client.DropletActions.Restore(id, image)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not restore droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Resize resizesx a droplet giving a size slug and optionally expands the disk.
func DropletActionResize(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	size := c.String(ArgImageSlug)
	disk := c.Bool(ArgResizeDisk)

	r, _, err := client.DropletActions.Resize(id, size, disk)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not resize droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Rebuild rebuilds a droplet using an image id or slug.
func DropletActionRebuild(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	image := c.String(ArgImage)

	var r *godo.Action
	var err error
	if i, aerr := strconv.Atoi(image); aerr == nil {
		fmt.Println("rebuilding by id")
		r, _, err = client.DropletActions.RebuildByImageID(id, i)
	} else {
		r, _, err = client.DropletActions.RebuildByImageSlug(id, image)
	}
	if err != nil {
		logrus.WithField("err", err).Fatal("could not rebuild droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Rename renames a droplet.
func DropletActionRename(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	name := c.String(ArgDropletName)

	r, _, err := client.DropletActions.Rename(id, name)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not rename droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// ChangeKernel changes the kernel for a droplet.
func DropletActionChangeKernel(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	kernel := c.Int(ArgKernelID)

	r, _, err := client.DropletActions.ChangeKernel(id, kernel)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not change droplet kernel")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Snapshot creates a snapshot for a droplet.
func DropletActionSnapshot(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	name := c.String(ArgSnapshotName)

	r, _, err := client.DropletActions.Snapshot(id, name)
	if err != nil {
		logrus.WithField("err", err).Fatal("could create snapshot for droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}
<commit_msg>removing unused debug statement<commit_after>package docli

import (
	"strconv"

	"github.com/Sirupsen/logrus"
	"github.com/codegangsta/cli"
	"github.com/digitalocean/godo"
)

// Get returns a droplet action by id.
func DropletActionGet(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	dropletID := c.Int(ArgDropletID)
	actionID := c.Int(ArgActionID)

	a, _, err := client.DropletActions.Get(dropletID, actionID)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not disable backups for droplet")
	}

	err = WriteJSON(a, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// DisableBackups disables backups for a droplet.
func DropletActionDisableBackups(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.DisableBackups(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not disable backups for droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Reboot reboots a droplet.
func DropletActionReboot(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.Reboot(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not reboot droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// PowerCycle power cycles a droplet.
func DropletActionPowerCycle(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	r, _, err := client.DropletActions.PowerCycle(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not power cycle droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Shutdown shuts a droplet down.
func DropletActionShutdown(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.Shutdown(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not shutdown droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// PowerOff turns droplet power off.
func DropletActionPowerOff(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.PowerOff(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not power off droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// PowerOn turns droplet power on.
func DropletActionPowerOn(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.PowerOn(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not power on droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// PasswordReset resets the droplet root password.
func DropletActionPasswordReset(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.PasswordReset(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not reset password for droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// EnableIPv6 enables IPv6 for a droplet.
func DropletActionEnableIPv6(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.EnableIPv6(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not enable IPv6 for droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// EnablePrivateNetworking enables private networking for a droplet.
func DropletActionEnablePrivateNetworking(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.EnablePrivateNetworking(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not enable private networking for droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Upgrade upgrades a droplet.
func DropletActionUpgrade(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)

	r, _, err := client.DropletActions.Upgrade(id)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not upgrade droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Restore restores a droplet using an image id.
func DropletActionRestore(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	image := c.Int(ArgImageID)

	r, _, err := client.DropletActions.Restore(id, image)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not restore droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Resize resizesx a droplet giving a size slug and optionally expands the disk.
func DropletActionResize(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	size := c.String(ArgImageSlug)
	disk := c.Bool(ArgResizeDisk)

	r, _, err := client.DropletActions.Resize(id, size, disk)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not resize droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Rebuild rebuilds a droplet using an image id or slug.
func DropletActionRebuild(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	image := c.String(ArgImage)

	var r *godo.Action
	var err error
	if i, aerr := strconv.Atoi(image); aerr == nil {
		r, _, err = client.DropletActions.RebuildByImageID(id, i)
	} else {
		r, _, err = client.DropletActions.RebuildByImageSlug(id, image)
	}
	if err != nil {
		logrus.WithField("err", err).Fatal("could not rebuild droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Rename renames a droplet.
func DropletActionRename(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	name := c.String(ArgDropletName)

	r, _, err := client.DropletActions.Rename(id, name)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not rename droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// ChangeKernel changes the kernel for a droplet.
func DropletActionChangeKernel(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	kernel := c.Int(ArgKernelID)

	r, _, err := client.DropletActions.ChangeKernel(id, kernel)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not change droplet kernel")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}

// Snapshot creates a snapshot for a droplet.
func DropletActionSnapshot(c *cli.Context) {
	client := NewClient(c, DefaultClientSource)
	id := c.Int(ArgDropletID)
	name := c.String(ArgSnapshotName)

	r, _, err := client.DropletActions.Snapshot(id, name)
	if err != nil {
		logrus.WithField("err", err).Fatal("could create snapshot for droplet")
	}

	err = WriteJSON(r, c.App.Writer)
	if err != nil {
		logrus.WithField("err", err).Fatal("could not write JSON")
	}
}
<|endoftext|>


<commit_before>// 文件数据记录类
// create by gloomy 2017-04-06 10:11:35
package common

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

const maxFileDataRecordingBytes = 1000000 // 默认文件大小

// 文件数据记录对象
// create by gloomy 2017-04-06 10:15:00
type FileDataRecording struct {
	sync.Mutex                         // 锁
	F                         *os.File // 文件对象
	FilePre                   string   // 文件开头字符串
	Fn                        string   // 文件路径
	Bytes                     int      // 文件大小
	Seq                       int      // 第几个
	FileProgram               string   // 文件存放路径
	MaxFileDataRecordingBytes int      // 文件大小
}

// 打开文件数据记录
// create by gloomy 2017-04-06 10:17:38
// 文件存放目录地址 文件开头字符串 文件大小
// 文件数据对象
func OpenLoadFile(fileProgram, filePre string, maxSize int) *FileDataRecording {
	if maxSize == 0 {
		maxSize = maxFileDataRecordingBytes
	}
	lf := &FileDataRecording{
		FilePre:                   filePre,
		FileProgram:               fileProgram,
		MaxFileDataRecordingBytes: maxSize,
	}
	lf.Rotate()
	return lf
}

// 文件退出
// create by gloomy 2017-04-06 10:27:58
func (f *FileDataRecording) Exit() {
	f.Lock()
	f.Close()
	f.F = nil
	f.Unlock()
}

// 文件关闭
// create by gloomy 2017-04-06 10:22:14
func (f *FileDataRecording) Close() {
	if f.F != nil {
		f.F.Close()
		os.Rename(f.Fn, f.Fn[0:len(f.Fn)-4]) //去掉末尾的.tmp
	}
}

// 文件切换
// create by gloomy 2017-04-06 10:30:05
func (f *FileDataRecording) Rotate() {
	f.Lock()
	f.Seq = 0
	f.Close()
	f.CreateNewFile()
	f.Unlock()
}

/// 创建新文件
// create by gloomy 2017-04-06 10:33:11
// 错误对象
func (f *FileDataRecording) CreateNewFile() (err error) {
	f.Bytes = 0
	if !strings.HasSuffix(f.FileProgram, "/") {
		f.FileProgram += "/"
	}
	f.Fn = fmt.Sprintf("%s%s-%d-%d.tmp", f.FileProgram, f.FilePre, time.Now().UnixNano(), f.Seq)
	f.Seq++
	f.F, err = os.OpenFile(f.Fn, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		fmt.Printf("create file %s failed: %s \n", f.Fn, err.Error())
	}
	return
}

// 写入数据
// create by gloomy 2017-04-06 11:40:55
// 需要写入的数据
// 错误对象
func (f *FileDataRecording) WriteData(dataStr string) (err error) {
	f.Lock()
	defer f.Unlock()
	if f.F == nil {
		err = f.CreateNewFile()
		if err != nil {
			return
		}
	}
	dataStrLen := len(dataStr)
	if f.Bytes+dataStrLen > f.MaxFileDataRecordingBytes {
		f.Close()
		if err = f.CreateNewFile(); err != nil {
			return
		}
	}
	f.Bytes += dataStrLen
	_, err = f.F.WriteString(dataStr)
	return
}

// 获取所有完成的文件列表
// create by gloomy 2017-04-06 13:46:51
// 文件列表
func (f *FileDataRecording) FileList() *[]string {
	var (
		fileArray      []string
		pathSplitArray []string
	)
	filepath.Walk(f.FileProgram, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}
		if ext := filepath.Ext(path); ext == ".tmp" {
			return nil
		}
		if info.Size() == 0 {
			os.Remove(path)
			return nil
		}
		pathSplitArray = strings.Split(path, "/")
		if strings.HasPrefix(pathSplitArray[len(pathSplitArray)-1], f.FilePre) {
			fileArray = append(fileArray, path)
		}
		return nil
	})
	return &fileArray
}
<commit_msg>删除没有作用的变量<commit_after>// 文件数据记录类
// create by gloomy 2017-04-06 10:11:35
package common

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"
)

const maxFileDataRecordingBytes = 1000000 // 默认文件大小

// 文件数据记录对象
// create by gloomy 2017-04-06 10:15:00
type FileDataRecording struct {
	sync.Mutex                         // 锁
	F                         *os.File // 文件对象
	FilePre                   string   // 文件开头字符串
	Fn                        string   // 文件路径
	Bytes                     int      // 文件大小
	Seq                       int      // 第几个
	FileProgram               string   // 文件存放路径
	MaxFileDataRecordingBytes int      // 文件大小
}

// 打开文件数据记录
// create by gloomy 2017-04-06 10:17:38
// 文件存放目录地址 文件开头字符串 文件大小
// 文件数据对象
func OpenLoadFile(fileProgram, filePre string, maxSize int) *FileDataRecording {
	if maxSize == 0 {
		maxSize = maxFileDataRecordingBytes
	}
	lf := &FileDataRecording{
		FilePre:                   filePre,
		FileProgram:               fileProgram,
		MaxFileDataRecordingBytes: maxSize,
	}
	lf.Rotate()
	return lf
}

// 文件退出
// create by gloomy 2017-04-06 10:27:58
func (f *FileDataRecording) Exit() {
	f.Lock()
	f.Close()
	f.F = nil
	f.Unlock()
}

// 文件关闭
// create by gloomy 2017-04-06 10:22:14
func (f *FileDataRecording) Close() {
	if f.F != nil {
		f.F.Close()
		os.Rename(f.Fn, f.Fn[0:len(f.Fn)-4]) //去掉末尾的.tmp
	}
}

// 文件切换
// create by gloomy 2017-04-06 10:30:05
func (f *FileDataRecording) Rotate() {
	f.Lock()
	f.Seq = 0
	f.Close()
	f.CreateNewFile()
	f.Unlock()
}

/// 创建新文件
// create by gloomy 2017-04-06 10:33:11
// 错误对象
func (f *FileDataRecording) CreateNewFile() (err error) {
	f.Bytes = 0
	if !strings.HasSuffix(f.FileProgram, "/") {
		f.FileProgram += "/"
	}
	f.Fn = fmt.Sprintf("%s%s-%d-%d.tmp", f.FileProgram, f.FilePre, time.Now().UnixNano(), f.Seq)
	f.Seq++
	f.F, err = os.OpenFile(f.Fn, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		fmt.Printf("create file %s failed: %s \n", f.Fn, err.Error())
	}
	return
}

// 写入数据
// create by gloomy 2017-04-06 11:40:55
// 需要写入的数据
// 错误对象
func (f *FileDataRecording) WriteData(dataStr string) (err error) {
	f.Lock()
	defer f.Unlock()
	if f.F == nil {
		err = f.CreateNewFile()
		if err != nil {
			return
		}
	}
	dataStrLen := len(dataStr)
	if f.Bytes+dataStrLen > f.MaxFileDataRecordingBytes {
		f.Close()
		if err = f.CreateNewFile(); err != nil {
			return
		}
	}
	f.Bytes += dataStrLen
	_, err = f.F.WriteString(dataStr)
	return
}

// 获取所有完成的文件列表
// create by gloomy 2017-04-06 13:46:51
// 文件列表
func (f *FileDataRecording) FileList() *[]string {
	var (
		fileArray []string
	)
	filepath.Walk(f.FileProgram, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}
		if ext := filepath.Ext(path); ext == ".tmp" {
			return nil
		}
		if info.Size() == 0 {
			os.Remove(path)
			return nil
		}
		if strings.HasPrefix(path, f.FilePre) {
			fileArray = append(fileArray, path)
		}
		return nil
	})
	return &fileArray
}
<|endoftext|>


<commit_before>package chainview

import (
	"fmt"
	"sync"
	"sync/atomic"

	"github.com/lightninglabs/neutrino"
	"github.com/roasbeef/btcd/chaincfg/chainhash"
	"github.com/roasbeef/btcd/rpcclient"
	"github.com/roasbeef/btcd/wire"
	"github.com/roasbeef/btcutil"
	"github.com/roasbeef/btcutil/gcs/builder"
	"github.com/roasbeef/btcwallet/waddrmgr"
)

// CfFilteredChainView is an implementation of the FilteredChainView interface
// which is supported by an underlying Bitcoin light client which supports
// client side filtering of Golomb Coded Sets. Rather than fetching all the
// blocks, the light client is able to query fitlers locally, to test if an
// item in a block modifies any of our watched set of UTXOs.
type CfFilteredChainView struct {
	started int32
	stopped int32

	// p2pNode is a pointer to the running GCS-filter supported Bitcoin
	// light clientl
	p2pNode *neutrino.ChainService

	// chainView is the active rescan which only watches our specified
	// sub-set of the UTXO set.
	chainView neutrino.Rescan

	// rescanErrChan is the channel that any errors encountered during the
	// rescan will be sent over.
	rescanErrChan <-chan error

	// newBlocks is the channel in which new filtered blocks are sent over.
	newBlocks chan *FilteredBlock

	// staleBlocks is the channel in which blocks that have been
	// disconnected from the mainchain are sent over.
	staleBlocks chan *FilteredBlock

	// filterUpdates is a channel in which updates to the utxo filter
	// attached to this instance are sent over.
	filterUpdates chan filterUpdate

	// chainFilter is the
	filterMtx   sync.RWMutex
	chainFilter map[wire.OutPoint]struct{}

	quit chan struct{}
	wg   sync.WaitGroup
}

// A compile time check to ensure CfFilteredChainView implements the
// chainview.FilteredChainView.
var _ FilteredChainView = (*CfFilteredChainView)(nil)

// NewCfFilteredChainView creates a new instance of the CfFilteredChainView
// which is connected to an active neutrino node.
//
// NOTE: The node should already be running an syncing before being passed into
// this function.
func NewCfFilteredChainView(node *neutrino.ChainService) (*CfFilteredChainView, error) {
	return &CfFilteredChainView{
		newBlocks:     make(chan *FilteredBlock),
		staleBlocks:   make(chan *FilteredBlock),
		quit:          make(chan struct{}),
		rescanErrChan: make(chan error),
		filterUpdates: make(chan filterUpdate),
		chainFilter:   make(map[wire.OutPoint]struct{}),
		p2pNode:       node,
	}, nil
}

// Start kicks off the FilteredChainView implementation. This function must be
// called before any calls to UpdateFilter can be processed.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) Start() error {
	// Already started?
	if atomic.AddInt32(&c.started, 1) != 1 {
		return nil
	}

	log.Infof("FilteredChainView starting")

	// First, we'll obtain the latest block height of the p2p node. We'll
	// start the auto-rescan from this point. Once a caller actually wishes
	// to register a chain view, the rescan state will be rewound
	// accordingly.
	bestHeader, bestHeight, err := c.p2pNode.BlockHeaders.ChainTip()
	if err != nil {
		return err
	}
	startingPoint := &waddrmgr.BlockStamp{
		Height: int32(bestHeight),
		Hash:   bestHeader.BlockHash(),
	}

	// Next, we'll create our set of rescan options. Currently it's
	// required that a user MUST set a addr/outpoint/txid when creating a
	// rescan. To get around this, we'll add a "zero" outpoint, that won't
	// actually be matched.
	var zeroPoint wire.OutPoint
	rescanOptions := []neutrino.RescanOption{
		neutrino.StartBlock(startingPoint),
		neutrino.QuitChan(c.quit),
		neutrino.NotificationHandlers(
			rpcclient.NotificationHandlers{
				OnFilteredBlockConnected:    c.onFilteredBlockConnected,
				OnFilteredBlockDisconnected: c.onFilteredBlockDisconnected,
			},
		),
		neutrino.WatchOutPoints(zeroPoint),
	}

	// Finally, we'll create our rescan struct, start it, and launch all
	// the goroutines we need to operate this FilteredChainView instance.
	c.chainView = c.p2pNode.NewRescan(rescanOptions...)
	c.rescanErrChan = c.chainView.Start()

	c.wg.Add(1)
	go c.chainFilterer()

	return nil
}

// Stop signals all active goroutines for a graceful shutdown.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) Stop() error {
	// Already shutting down?
	if atomic.AddInt32(&c.stopped, 1) != 1 {
		return nil
	}

	log.Infof("FilteredChainView stopping")

	close(c.quit)
	c.wg.Wait()

	return nil
}

// onFilteredBlockConnected is called for each block that's connected to the
// end of the main chain. Based on our current chain filter, the block may or
// may not include any relevant transactions.
func (c *CfFilteredChainView) onFilteredBlockConnected(height int32,
	header *wire.BlockHeader, txns []*btcutil.Tx) {

	mtxs := make([]*wire.MsgTx, len(txns))
	for i, tx := range txns {
		mtx := tx.MsgTx()
		mtxs[i] = mtx

		for _, txIn := range mtx.TxIn {
			c.filterMtx.Lock()
			delete(c.chainFilter, txIn.PreviousOutPoint)
			c.filterMtx.Unlock()
		}

	}

	go func() {
		c.newBlocks <- &FilteredBlock{
			Hash:         header.BlockHash(),
			Height:       uint32(height),
			Transactions: mtxs,
		}
	}()
}

// onFilteredBlockDisconnected is a callback which is executed once a block is
// disconnected from the end of the main chain.
func (c *CfFilteredChainView) onFilteredBlockDisconnected(height int32,
	header *wire.BlockHeader) {

	go func() {
		c.staleBlocks <- &FilteredBlock{
			Hash:   header.BlockHash(),
			Height: uint32(height),
		}
	}()
}

// chainFilterer is the primary coordination goroutine within the
// CfFilteredChainView. This goroutine handles errors from the running rescan,
// and also filter updates.
func (c *CfFilteredChainView) chainFilterer() {
	defer c.wg.Done()

	for {
		select {

		case err := <-c.rescanErrChan:
			log.Errorf("Error encountered during rescan: %v", err)

		// We've received a new update to the filter from the caller to
		// mutate their established chain view.
		case update := <-c.filterUpdates:
			log.Debugf("Updating chain filter with new UTXO's: %v",
				update.newUtxos)

			// First, we'll update the current chain view, by
			// adding any new UTXO's, ignoring duplicates int he
			// process.
			c.filterMtx.Lock()
			for _, op := range update.newUtxos {
				c.chainFilter[op] = struct{}{}
			}
			c.filterMtx.Unlock()

			// With our internal chain view update, we'll craft a
			// new update to the chainView which includes our new
			// UTXO's, and current update height.
			rescanUpdate := []neutrino.UpdateOption{
				neutrino.AddOutPoints(update.newUtxos...),
				neutrino.Rewind(update.updateHeight),
			}
			err := c.chainView.Update(rescanUpdate...)
			if err != nil {
				log.Errorf("unable to update rescan: %v", err)
			}

			if update.done != nil {
				close(update.done)
			}

		case <-c.quit:
			return
		}
	}
}

// FilterBlock takes a block hash, and returns a FilteredBlocks which is the
// result of applying the current registered UTXO sub-set on the block
// corresponding to that block hash. If any watched UTXO's are spent by the
// selected lock, then the internal chainFilter will also be updated.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) FilterBlock(blockHash *chainhash.Hash) (*FilteredBlock, error) {
	// First, we'll fetch the block header itself so we can obtain the
	// height which is part of our return value.
	_, blockHeight, err := c.p2pNode.BlockHeaders.FetchHeader(blockHash)
	if err != nil {
		return nil, err
	}

	filteredBlock := &FilteredBlock{
		Hash:   *blockHash,
		Height: blockHeight,
	}

	// If we don't have any items within our current chain filter, then we
	// can exit early as we don't need to fetch the filter.
	c.filterMtx.RLock()
	if len(c.chainFilter) == 0 {
		c.filterMtx.RUnlock()
		return filteredBlock, nil
	}
	c.filterMtx.RUnlock()

	// Next, using the block, hash, we'll fetch the compact filter for this
	// block. We only require the regular filter as we're just looking for
	// outpoint that have been spent.
	filter, err := c.p2pNode.GetCFilter(*blockHash, false)
	if err != nil {
		return nil, err
	}

	// Before we can match the filter, we'll need to map each item in our
	// chain filter to the representation that included in the compact
	// filters.
	c.filterMtx.RLock()
	relevantPoints := make([][]byte, 0, len(c.chainFilter))
	for op := range c.chainFilter {
		opBytes := builder.OutPointToFilterEntry(op)
		relevantPoints = append(relevantPoints, opBytes)
	}
	c.filterMtx.RUnlock()

	// With our relevant points constructed, we can finally match against
	// the retrieved filter.
	matched, err := filter.MatchAny(builder.DeriveKey(blockHash),
		relevantPoints)
	if err != nil {
		return nil, err
	}

	// If there wasn't a match, then we'll return the filtered block as is
	// (void of any transactions).
	if !matched {
		return filteredBlock, nil
	}

	// If we reach this point, then there was a match, so we'll need to
	// fetch the block itself so we can scan it for any actual matches (as
	// there's a fp rate).
	block, err := c.p2pNode.GetBlockFromNetwork(*blockHash)
	if err != nil {
		return nil, err
	}

	// Finally, we'll step through the block, input by input, to see if any
	// transactions spend any outputs from our watched sub-set of the UTXO
	// set.
	for _, tx := range block.Transactions() {
		for _, txIn := range tx.MsgTx().TxIn {
			prevOp := txIn.PreviousOutPoint

			c.filterMtx.RLock()
			_, ok := c.chainFilter[prevOp]
			c.filterMtx.RUnlock()

			if ok {
				filteredBlock.Transactions = append(
					filteredBlock.Transactions,
					tx.MsgTx(),
				)

				c.filterMtx.Lock()
				delete(c.chainFilter, prevOp)
				c.filterMtx.Unlock()

				break
			}
		}
	}

	return filteredBlock, nil
}

// UpdateFilter updates the UTXO filter which is to be consulted when creating
// FilteredBlocks to be sent to subscribed clients. This method is cumulative
// meaning repeated calls to this method should _expand_ the size of the UTXO
// sub-set currently being watched.  If the set updateHeight is _lower_ than
// the best known height of the implementation, then the state should be
// rewound to ensure all relevant notifications are dispatched.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) UpdateFilter(ops []wire.OutPoint, updateHeight uint32) error {
	doneChan := make(chan struct{})
	update := filterUpdate{
		newUtxos:     ops,
		updateHeight: updateHeight,
		done:         doneChan,
	}

	select {
	case c.filterUpdates <- update:
	case <-c.quit:
		return fmt.Errorf("chain filter shutting down")
	}

	select {
	case <-doneChan:
		return nil
	case <-c.quit:
		return fmt.Errorf("chain filter shutting down")
	}

}

// FilteredBlocks returns the channel that filtered blocks are to be sent over.
// Each time a block is connected to the end of a main chain, and appropriate
// FilteredBlock which contains the transactions which mutate our watched UTXO
// set is to be returned.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) FilteredBlocks() <-chan *FilteredBlock {
	return c.newBlocks
}

// DisconnectedBlocks returns a receive only channel which will be sent upon
// with the empty filtered blocks of blocks which are disconnected from the
// main chain in the case of a re-org.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) DisconnectedBlocks() <-chan *FilteredBlock {
	return c.staleBlocks
}
<commit_msg>routing/chainview: Fix data race in block disconnected callback.<commit_after>package chainview

import (
	"fmt"
	"sync"
	"sync/atomic"

	"github.com/lightninglabs/neutrino"
	"github.com/roasbeef/btcd/chaincfg/chainhash"
	"github.com/roasbeef/btcd/rpcclient"
	"github.com/roasbeef/btcd/wire"
	"github.com/roasbeef/btcutil"
	"github.com/roasbeef/btcutil/gcs/builder"
	"github.com/roasbeef/btcwallet/waddrmgr"
)

// CfFilteredChainView is an implementation of the FilteredChainView interface
// which is supported by an underlying Bitcoin light client which supports
// client side filtering of Golomb Coded Sets. Rather than fetching all the
// blocks, the light client is able to query fitlers locally, to test if an
// item in a block modifies any of our watched set of UTXOs.
type CfFilteredChainView struct {
	started int32
	stopped int32

	// p2pNode is a pointer to the running GCS-filter supported Bitcoin
	// light clientl
	p2pNode *neutrino.ChainService

	// chainView is the active rescan which only watches our specified
	// sub-set of the UTXO set.
	chainView neutrino.Rescan

	// rescanErrChan is the channel that any errors encountered during the
	// rescan will be sent over.
	rescanErrChan <-chan error

	// newBlocks is the channel in which new filtered blocks are sent over.
	newBlocks chan *FilteredBlock

	// staleBlocks is the channel in which blocks that have been
	// disconnected from the mainchain are sent over.
	staleBlocks chan *FilteredBlock

	// filterUpdates is a channel in which updates to the utxo filter
	// attached to this instance are sent over.
	filterUpdates chan filterUpdate

	// chainFilter is the
	filterMtx   sync.RWMutex
	chainFilter map[wire.OutPoint]struct{}

	quit chan struct{}
	wg   sync.WaitGroup
}

// A compile time check to ensure CfFilteredChainView implements the
// chainview.FilteredChainView.
var _ FilteredChainView = (*CfFilteredChainView)(nil)

// NewCfFilteredChainView creates a new instance of the CfFilteredChainView
// which is connected to an active neutrino node.
//
// NOTE: The node should already be running an syncing before being passed into
// this function.
func NewCfFilteredChainView(node *neutrino.ChainService) (*CfFilteredChainView, error) {
	return &CfFilteredChainView{
		newBlocks:     make(chan *FilteredBlock),
		staleBlocks:   make(chan *FilteredBlock),
		quit:          make(chan struct{}),
		rescanErrChan: make(chan error),
		filterUpdates: make(chan filterUpdate),
		chainFilter:   make(map[wire.OutPoint]struct{}),
		p2pNode:       node,
	}, nil
}

// Start kicks off the FilteredChainView implementation. This function must be
// called before any calls to UpdateFilter can be processed.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) Start() error {
	// Already started?
	if atomic.AddInt32(&c.started, 1) != 1 {
		return nil
	}

	log.Infof("FilteredChainView starting")

	// First, we'll obtain the latest block height of the p2p node. We'll
	// start the auto-rescan from this point. Once a caller actually wishes
	// to register a chain view, the rescan state will be rewound
	// accordingly.
	bestHeader, bestHeight, err := c.p2pNode.BlockHeaders.ChainTip()
	if err != nil {
		return err
	}
	startingPoint := &waddrmgr.BlockStamp{
		Height: int32(bestHeight),
		Hash:   bestHeader.BlockHash(),
	}

	// Next, we'll create our set of rescan options. Currently it's
	// required that a user MUST set a addr/outpoint/txid when creating a
	// rescan. To get around this, we'll add a "zero" outpoint, that won't
	// actually be matched.
	var zeroPoint wire.OutPoint
	rescanOptions := []neutrino.RescanOption{
		neutrino.StartBlock(startingPoint),
		neutrino.QuitChan(c.quit),
		neutrino.NotificationHandlers(
			rpcclient.NotificationHandlers{
				OnFilteredBlockConnected:    c.onFilteredBlockConnected,
				OnFilteredBlockDisconnected: c.onFilteredBlockDisconnected,
			},
		),
		neutrino.WatchOutPoints(zeroPoint),
	}

	// Finally, we'll create our rescan struct, start it, and launch all
	// the goroutines we need to operate this FilteredChainView instance.
	c.chainView = c.p2pNode.NewRescan(rescanOptions...)
	c.rescanErrChan = c.chainView.Start()

	c.wg.Add(1)
	go c.chainFilterer()

	return nil
}

// Stop signals all active goroutines for a graceful shutdown.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) Stop() error {
	// Already shutting down?
	if atomic.AddInt32(&c.stopped, 1) != 1 {
		return nil
	}

	log.Infof("FilteredChainView stopping")

	close(c.quit)
	c.wg.Wait()

	return nil
}

// onFilteredBlockConnected is called for each block that's connected to the
// end of the main chain. Based on our current chain filter, the block may or
// may not include any relevant transactions.
func (c *CfFilteredChainView) onFilteredBlockConnected(height int32,
	header *wire.BlockHeader, txns []*btcutil.Tx) {

	mtxs := make([]*wire.MsgTx, len(txns))
	for i, tx := range txns {
		mtx := tx.MsgTx()
		mtxs[i] = mtx

		for _, txIn := range mtx.TxIn {
			c.filterMtx.Lock()
			delete(c.chainFilter, txIn.PreviousOutPoint)
			c.filterMtx.Unlock()
		}

	}

	go func() {
		c.newBlocks <- &FilteredBlock{
			Hash:         header.BlockHash(),
			Height:       uint32(height),
			Transactions: mtxs,
		}
	}()
}

// onFilteredBlockDisconnected is a callback which is executed once a block is
// disconnected from the end of the main chain.
func (c *CfFilteredChainView) onFilteredBlockDisconnected(height int32,
	header *wire.BlockHeader) {

	filteredBlock := &FilteredBlock{
		Hash:   header.BlockHash(),
		Height: uint32(height),
	}

	go func() {
		c.staleBlocks <- filteredBlock
	}()
}

// chainFilterer is the primary coordination goroutine within the
// CfFilteredChainView. This goroutine handles errors from the running rescan,
// and also filter updates.
func (c *CfFilteredChainView) chainFilterer() {
	defer c.wg.Done()

	for {
		select {

		case err := <-c.rescanErrChan:
			log.Errorf("Error encountered during rescan: %v", err)

		// We've received a new update to the filter from the caller to
		// mutate their established chain view.
		case update := <-c.filterUpdates:
			log.Debugf("Updating chain filter with new UTXO's: %v",
				update.newUtxos)

			// First, we'll update the current chain view, by
			// adding any new UTXO's, ignoring duplicates int he
			// process.
			c.filterMtx.Lock()
			for _, op := range update.newUtxos {
				c.chainFilter[op] = struct{}{}
			}
			c.filterMtx.Unlock()

			// With our internal chain view update, we'll craft a
			// new update to the chainView which includes our new
			// UTXO's, and current update height.
			rescanUpdate := []neutrino.UpdateOption{
				neutrino.AddOutPoints(update.newUtxos...),
				neutrino.Rewind(update.updateHeight),
			}
			err := c.chainView.Update(rescanUpdate...)
			if err != nil {
				log.Errorf("unable to update rescan: %v", err)
			}

			if update.done != nil {
				close(update.done)
			}

		case <-c.quit:
			return
		}
	}
}

// FilterBlock takes a block hash, and returns a FilteredBlocks which is the
// result of applying the current registered UTXO sub-set on the block
// corresponding to that block hash. If any watched UTXO's are spent by the
// selected lock, then the internal chainFilter will also be updated.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) FilterBlock(blockHash *chainhash.Hash) (*FilteredBlock, error) {
	// First, we'll fetch the block header itself so we can obtain the
	// height which is part of our return value.
	_, blockHeight, err := c.p2pNode.BlockHeaders.FetchHeader(blockHash)
	if err != nil {
		return nil, err
	}

	filteredBlock := &FilteredBlock{
		Hash:   *blockHash,
		Height: blockHeight,
	}

	// If we don't have any items within our current chain filter, then we
	// can exit early as we don't need to fetch the filter.
	c.filterMtx.RLock()
	if len(c.chainFilter) == 0 {
		c.filterMtx.RUnlock()
		return filteredBlock, nil
	}
	c.filterMtx.RUnlock()

	// Next, using the block, hash, we'll fetch the compact filter for this
	// block. We only require the regular filter as we're just looking for
	// outpoint that have been spent.
	filter, err := c.p2pNode.GetCFilter(*blockHash, false)
	if err != nil {
		return nil, err
	}

	// Before we can match the filter, we'll need to map each item in our
	// chain filter to the representation that included in the compact
	// filters.
	c.filterMtx.RLock()
	relevantPoints := make([][]byte, 0, len(c.chainFilter))
	for op := range c.chainFilter {
		opBytes := builder.OutPointToFilterEntry(op)
		relevantPoints = append(relevantPoints, opBytes)
	}
	c.filterMtx.RUnlock()

	// With our relevant points constructed, we can finally match against
	// the retrieved filter.
	matched, err := filter.MatchAny(builder.DeriveKey(blockHash),
		relevantPoints)
	if err != nil {
		return nil, err
	}

	// If there wasn't a match, then we'll return the filtered block as is
	// (void of any transactions).
	if !matched {
		return filteredBlock, nil
	}

	// If we reach this point, then there was a match, so we'll need to
	// fetch the block itself so we can scan it for any actual matches (as
	// there's a fp rate).
	block, err := c.p2pNode.GetBlockFromNetwork(*blockHash)
	if err != nil {
		return nil, err
	}

	// Finally, we'll step through the block, input by input, to see if any
	// transactions spend any outputs from our watched sub-set of the UTXO
	// set.
	for _, tx := range block.Transactions() {
		for _, txIn := range tx.MsgTx().TxIn {
			prevOp := txIn.PreviousOutPoint

			c.filterMtx.RLock()
			_, ok := c.chainFilter[prevOp]
			c.filterMtx.RUnlock()

			if ok {
				filteredBlock.Transactions = append(
					filteredBlock.Transactions,
					tx.MsgTx(),
				)

				c.filterMtx.Lock()
				delete(c.chainFilter, prevOp)
				c.filterMtx.Unlock()

				break
			}
		}
	}

	return filteredBlock, nil
}

// UpdateFilter updates the UTXO filter which is to be consulted when creating
// FilteredBlocks to be sent to subscribed clients. This method is cumulative
// meaning repeated calls to this method should _expand_ the size of the UTXO
// sub-set currently being watched.  If the set updateHeight is _lower_ than
// the best known height of the implementation, then the state should be
// rewound to ensure all relevant notifications are dispatched.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) UpdateFilter(ops []wire.OutPoint, updateHeight uint32) error {
	doneChan := make(chan struct{})
	update := filterUpdate{
		newUtxos:     ops,
		updateHeight: updateHeight,
		done:         doneChan,
	}

	select {
	case c.filterUpdates <- update:
	case <-c.quit:
		return fmt.Errorf("chain filter shutting down")
	}

	select {
	case <-doneChan:
		return nil
	case <-c.quit:
		return fmt.Errorf("chain filter shutting down")
	}

}

// FilteredBlocks returns the channel that filtered blocks are to be sent over.
// Each time a block is connected to the end of a main chain, and appropriate
// FilteredBlock which contains the transactions which mutate our watched UTXO
// set is to be returned.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) FilteredBlocks() <-chan *FilteredBlock {
	return c.newBlocks
}

// DisconnectedBlocks returns a receive only channel which will be sent upon
// with the empty filtered blocks of blocks which are disconnected from the
// main chain in the case of a re-org.
//
// NOTE: This is part of the FilteredChainView interface.
func (c *CfFilteredChainView) DisconnectedBlocks() <-chan *FilteredBlock {
	return c.staleBlocks
}
<|endoftext|>


<commit_before>package varint

import (
	"io"
)

// ReadFrom reads a varint encoded from the given io.Reader.
//
// On success, the function returns the varint as a int64, the number of bytes
// consumed, and nil.
func ReadFrom(r io.Reader) (int64, int64, error) {
	var buffer [3]byte
	if b, err := ReadByte(r); err != nil {
		return 0, 0, err
	} else {
		if (b & 0x80) == 0 {
			return int64(b), 1, nil
		}
		if (b & 0xC0) == 0x80 {
			if n, err := io.ReadFull(r, buffer[0:1]); err != nil {
				return 0, int64(n + 1), err
			}
			return int64(b&0x3F)<<8 | int64(buffer[0]), 2, nil
		}
		if (b & 0xE0) == 0xC0 {
			if n, err := io.ReadFull(r, buffer[0:2]); err != nil {
				return 0, int64(n + 1), err
			}
			return int64(b&0x1F)<<8 | int64(buffer[0])<<8 | int64(buffer[1]), 3, nil
		}
		if (b & 0xF0) == 0xE0 {
			if n, err := io.ReadFull(r, buffer[0:3]); err != nil {
				return 0, int64(n + 1), err
			}
			return int64(b&0xF)<<8 | int64(buffer[0])<<8 | int64(buffer[1])<<8 | int64(buffer[2]), 4, nil
		}
	}
	return 0, 0, ErrOutOfRange
}

// ReadBytes reads a single byte from the given io.Reader.
func ReadByte(r io.Reader) (byte, error) {
	var buff [1]byte
	if _, err := io.ReadFull(r, buff[:]); err != nil {
		return 0, err
	} else {
		return buff[0], nil
	}
}
<commit_msg>fix varint.ReadFrom return value<commit_after>package varint

import (
	"io"
)

// ReadFrom reads a varint encoded from the given io.Reader.
//
// On success, the function returns the varint as a int64, the number of bytes
// consumed, and nil.
func ReadFrom(r io.Reader) (int64, int64, error) {
	var buffer [3]byte
	if b, err := ReadByte(r); err != nil {
		return 0, 0, err
	} else {
		if (b & 0x80) == 0 {
			return int64(b), 1, nil
		}
		if (b & 0xC0) == 0x80 {
			if n, err := io.ReadFull(r, buffer[0:1]); err != nil {
				return 0, int64(n + 1), err
			}
			return int64(b&0x3F)<<8 | int64(buffer[0]), 2, nil
		}
		if (b & 0xE0) == 0xC0 {
			if n, err := io.ReadFull(r, buffer[0:2]); err != nil {
				return 0, int64(n + 1), err
			}
			return int64(b&0x1F)<<8 | int64(buffer[0])<<8 | int64(buffer[1]), 3, nil
		}
		if (b & 0xF0) == 0xE0 {
			if n, err := io.ReadFull(r, buffer[0:3]); err != nil {
				return 0, int64(n + 1), err
			}
			return int64(b&0xF)<<8 | int64(buffer[0])<<8 | int64(buffer[1])<<8 | int64(buffer[2]), 4, nil
		}
	}
	return 0, 1, ErrOutOfRange
}

// ReadBytes reads a single byte from the given io.Reader.
func ReadByte(r io.Reader) (byte, error) {
	var buff [1]byte
	if _, err := io.ReadFull(r, buff[:]); err != nil {
		return 0, err
	} else {
		return buff[0], nil
	}
}
<|endoftext|>


<commit_before>package cachememdb

import (
	"errors"
	"fmt"

	memdb "github.com/hashicorp/go-memdb"
)

const (
	tableNameIndexer = "indexer"
)

// CacheMemDB is the underlying cache database for storing indexes.
type CacheMemDB struct {
	db *memdb.MemDB
}

// New creates a new instance of CacheMemDB.
func New() (*CacheMemDB, error) {
	db, err := newDB()
	if err != nil {
		return nil, err
	}

	return &CacheMemDB{
		db: db,
	}, nil
}

func newDB() (*memdb.MemDB, error) {
	cacheSchema := &memdb.DBSchema{
		Tables: map[string]*memdb.TableSchema{
			tableNameIndexer: &memdb.TableSchema{
				Name: tableNameIndexer,
				Indexes: map[string]*memdb.IndexSchema{
					// This index enables fetching the cached item based on the
					// identifier of the index.
					IndexNameID: &memdb.IndexSchema{
						Name:   IndexNameID,
						Unique: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "ID",
						},
					},
					// This index enables fetching all the entries in cache for
					// a given request path, in a given namespace.
					IndexNameRequestPath: &memdb.IndexSchema{
						Name:   IndexNameRequestPath,
						Unique: false,
						Indexer: &memdb.CompoundIndex{
							Indexes: []memdb.Indexer{
								&memdb.StringFieldIndex{
									Field: "Namespace",
								},
								&memdb.StringFieldIndex{
									Field: "RequestPath",
								},
							},
						},
					},
					// This index enables fetching all the entries in cache
					// belonging to the leases of a given token.
					IndexNameLeaseToken: &memdb.IndexSchema{
						Name:         IndexNameLeaseToken,
						Unique:       false,
						AllowMissing: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "LeaseToken",
						},
					},
					// This index enables fetching all the entries in cache
					// that are tied to the given token, regardless of the
					// entries belonging to the token or belonging to the
					// lease.
					IndexNameToken: &memdb.IndexSchema{
						Name:         IndexNameToken,
						Unique:       true,
						AllowMissing: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "Token",
						},
					},
					// This index enables fetching all the entries in cache for
					// the given parent token.
					IndexNameTokenParent: &memdb.IndexSchema{
						Name:         IndexNameTokenParent,
						Unique:       false,
						AllowMissing: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "TokenParent",
						},
					},
					// This index enables fetching all the entries in cache for
					// the given accessor.
					IndexNameTokenAccessor: &memdb.IndexSchema{
						Name:         IndexNameTokenAccessor,
						Unique:       true,
						AllowMissing: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "TokenAccessor",
						},
					},
					// This index enables fetching all the entries in cache for
					// the given lease identifier.
					IndexNameLease: &memdb.IndexSchema{
						Name:         IndexNameLease,
						Unique:       true,
						AllowMissing: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "Lease",
						},
					},
				},
			},
		},
	}

	db, err := memdb.NewMemDB(cacheSchema)
	if err != nil {
		return nil, err
	}
	return db, nil
}

// Get returns the index based on the indexer and the index values provided.
func (c *CacheMemDB) Get(indexName string, indexValues ...interface{}) (*Index, error) {
	if !validIndexName(indexName) {
		return nil, fmt.Errorf("invalid index name %q", indexName)
	}

	raw, err := c.db.Txn(false).First(tableNameIndexer, indexName, indexValues...)
	if err != nil {
		return nil, err
	}

	if raw == nil {
		return nil, nil
	}

	index, ok := raw.(*Index)
	if !ok {
		return nil, errors.New("unable to parse index value from the cache")
	}

	return index, nil
}

// Set stores the index into the cache.
func (c *CacheMemDB) Set(index *Index) error {
	if index == nil {
		return errors.New("nil index provided")
	}

	txn := c.db.Txn(true)
	defer txn.Abort()

	if err := txn.Insert(tableNameIndexer, index); err != nil {
		return fmt.Errorf("unable to insert index into cache: %v", err)
	}

	txn.Commit()

	return nil
}

// GetByPrefix returns all the cached indexes based on the index name and the
// value prefix.
func (c *CacheMemDB) GetByPrefix(indexName string, indexValues ...interface{}) ([]*Index, error) {
	if !validIndexName(indexName) {
		return nil, fmt.Errorf("invalid index name %q", indexName)
	}

	indexName = indexName + "_prefix"

	// Get all the objects
	iter, err := c.db.Txn(false).Get(tableNameIndexer, indexName, indexValues...)
	if err != nil {
		return nil, err
	}

	var indexes []*Index
	for {
		obj := iter.Next()
		if obj == nil {
			break
		}
		index, ok := obj.(*Index)
		if !ok {
			return nil, fmt.Errorf("failed to cast cached index")
		}

		indexes = append(indexes, index)
	}

	return indexes, nil
}

// Evict removes an index from the cache based on index name and value.
func (c *CacheMemDB) Evict(indexName string, indexValues ...interface{}) error {
	index, err := c.Get(indexName, indexValues...)
	if err != nil {
		return fmt.Errorf("unable to fetch index on cache deletion: %v", err)
	}

	if index == nil {
		return nil
	}

	txn := c.db.Txn(true)
	defer txn.Abort()

	if err := txn.Delete(tableNameIndexer, index); err != nil {
		return fmt.Errorf("unable to delete index from cache: %v", err)
	}

	txn.Commit()

	return nil
}

// EvictAll removes all matching indexes from the cache based on index name and value.
func (c *CacheMemDB) EvictAll(indexName, indexValue string) error {
	return c.batchEvict(false, indexName, indexValue)
}

// EvictByPrefix removes all matching prefix indexes from the cache based on index name and prefix.
func (c *CacheMemDB) EvictByPrefix(indexName, indexPrefix string) error {
	return c.batchEvict(true, indexName, indexPrefix)
}

// batchEvict is a helper that supports eviction based on absolute and prefixed index values.
func (c *CacheMemDB) batchEvict(isPrefix bool, indexName string, indexValues ...interface{}) error {
	if !validIndexName(indexName) {
		return fmt.Errorf("invalid index name %q", indexName)
	}

	if isPrefix {
		indexName = indexName + "_prefix"
	}

	txn := c.db.Txn(true)
	defer txn.Abort()

	_, err := txn.DeleteAll(tableNameIndexer, indexName, indexValues...)
	if err != nil {
		return err
	}

	txn.Commit()

	return nil
}

// Flush resets the underlying cache object.
func (c *CacheMemDB) Flush() error {
	newDB, err := newDB()
	if err != nil {
		return err
	}

	c.db = newDB

	return nil
}
<commit_msg>agent/caching: remove EvictByPrefix and EvictAll (#6269)<commit_after>package cachememdb

import (
	"errors"
	"fmt"

	memdb "github.com/hashicorp/go-memdb"
)

const (
	tableNameIndexer = "indexer"
)

// CacheMemDB is the underlying cache database for storing indexes.
type CacheMemDB struct {
	db *memdb.MemDB
}

// New creates a new instance of CacheMemDB.
func New() (*CacheMemDB, error) {
	db, err := newDB()
	if err != nil {
		return nil, err
	}

	return &CacheMemDB{
		db: db,
	}, nil
}

func newDB() (*memdb.MemDB, error) {
	cacheSchema := &memdb.DBSchema{
		Tables: map[string]*memdb.TableSchema{
			tableNameIndexer: &memdb.TableSchema{
				Name: tableNameIndexer,
				Indexes: map[string]*memdb.IndexSchema{
					// This index enables fetching the cached item based on the
					// identifier of the index.
					IndexNameID: &memdb.IndexSchema{
						Name:   IndexNameID,
						Unique: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "ID",
						},
					},
					// This index enables fetching all the entries in cache for
					// a given request path, in a given namespace.
					IndexNameRequestPath: &memdb.IndexSchema{
						Name:   IndexNameRequestPath,
						Unique: false,
						Indexer: &memdb.CompoundIndex{
							Indexes: []memdb.Indexer{
								&memdb.StringFieldIndex{
									Field: "Namespace",
								},
								&memdb.StringFieldIndex{
									Field: "RequestPath",
								},
							},
						},
					},
					// This index enables fetching all the entries in cache
					// belonging to the leases of a given token.
					IndexNameLeaseToken: &memdb.IndexSchema{
						Name:         IndexNameLeaseToken,
						Unique:       false,
						AllowMissing: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "LeaseToken",
						},
					},
					// This index enables fetching all the entries in cache
					// that are tied to the given token, regardless of the
					// entries belonging to the token or belonging to the
					// lease.
					IndexNameToken: &memdb.IndexSchema{
						Name:         IndexNameToken,
						Unique:       true,
						AllowMissing: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "Token",
						},
					},
					// This index enables fetching all the entries in cache for
					// the given parent token.
					IndexNameTokenParent: &memdb.IndexSchema{
						Name:         IndexNameTokenParent,
						Unique:       false,
						AllowMissing: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "TokenParent",
						},
					},
					// This index enables fetching all the entries in cache for
					// the given accessor.
					IndexNameTokenAccessor: &memdb.IndexSchema{
						Name:         IndexNameTokenAccessor,
						Unique:       true,
						AllowMissing: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "TokenAccessor",
						},
					},
					// This index enables fetching all the entries in cache for
					// the given lease identifier.
					IndexNameLease: &memdb.IndexSchema{
						Name:         IndexNameLease,
						Unique:       true,
						AllowMissing: true,
						Indexer: &memdb.StringFieldIndex{
							Field: "Lease",
						},
					},
				},
			},
		},
	}

	db, err := memdb.NewMemDB(cacheSchema)
	if err != nil {
		return nil, err
	}
	return db, nil
}

// Get returns the index based on the indexer and the index values provided.
func (c *CacheMemDB) Get(indexName string, indexValues ...interface{}) (*Index, error) {
	if !validIndexName(indexName) {
		return nil, fmt.Errorf("invalid index name %q", indexName)
	}

	raw, err := c.db.Txn(false).First(tableNameIndexer, indexName, indexValues...)
	if err != nil {
		return nil, err
	}

	if raw == nil {
		return nil, nil
	}

	index, ok := raw.(*Index)
	if !ok {
		return nil, errors.New("unable to parse index value from the cache")
	}

	return index, nil
}

// Set stores the index into the cache.
func (c *CacheMemDB) Set(index *Index) error {
	if index == nil {
		return errors.New("nil index provided")
	}

	txn := c.db.Txn(true)
	defer txn.Abort()

	if err := txn.Insert(tableNameIndexer, index); err != nil {
		return fmt.Errorf("unable to insert index into cache: %v", err)
	}

	txn.Commit()

	return nil
}

// GetByPrefix returns all the cached indexes based on the index name and the
// value prefix.
func (c *CacheMemDB) GetByPrefix(indexName string, indexValues ...interface{}) ([]*Index, error) {
	if !validIndexName(indexName) {
		return nil, fmt.Errorf("invalid index name %q", indexName)
	}

	indexName = indexName + "_prefix"

	// Get all the objects
	iter, err := c.db.Txn(false).Get(tableNameIndexer, indexName, indexValues...)
	if err != nil {
		return nil, err
	}

	var indexes []*Index
	for {
		obj := iter.Next()
		if obj == nil {
			break
		}
		index, ok := obj.(*Index)
		if !ok {
			return nil, fmt.Errorf("failed to cast cached index")
		}

		indexes = append(indexes, index)
	}

	return indexes, nil
}

// Evict removes an index from the cache based on index name and value.
func (c *CacheMemDB) Evict(indexName string, indexValues ...interface{}) error {
	index, err := c.Get(indexName, indexValues...)
	if err != nil {
		return fmt.Errorf("unable to fetch index on cache deletion: %v", err)
	}

	if index == nil {
		return nil
	}

	txn := c.db.Txn(true)
	defer txn.Abort()

	if err := txn.Delete(tableNameIndexer, index); err != nil {
		return fmt.Errorf("unable to delete index from cache: %v", err)
	}

	txn.Commit()

	return nil
}

// Flush resets the underlying cache object.
func (c *CacheMemDB) Flush() error {
	newDB, err := newDB()
	if err != nil {
		return err
	}

	c.db = newDB

	return nil
}
<|endoftext|>