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

package monitor

import (
	"atlantis/supervisor/containers/serialize"
	"atlantis/supervisor/rpc/types"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/jigish/go-flags"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"
)

const (
	OK = iota
	Warning
	Critical
	Uknown
)

type Config struct {
	ContainerFile   string `toml:"container_file"`
	ContainersDir   string `toml:"container_dir"`
	InventoryDir    string `toml:"inventory_dir"`
	SSHIdentity     string `toml:"ssh_identity"`
	SSHUser         string `toml:"ssh_user"`
	CheckName       string `toml:"check_name"`
	CheckDir        string `toml:"check_dir"`
	DefaultGroup    string `toml:"default_group"`
	TimeoutDuration uint   `toml:"timeout_duration"`
	Verbose         bool   `toml:"verbose"`
}

type Opts struct {
	ContainerFile   string `short:"f" long:"container-file" description:"file to get container information"`
	ContainersDir   string `short:"s" long:"containers-dir" description:"directory containing configs for each container"`
	SSHIdentity     string `short:"i" long:"ssh-identity" description:"file containing the SSH key for all containers"`
	SSHUser         string `short:"u" long:"ssh-user" description:"user account to ssh into containers"`
	CheckName       string `short:"n" long:"check-name" description:"service name that will appear in Nagios for the monitor"`
	CheckDir        string `short:"d" long:"check-dir" description:"directory containing all the scripts for the monitoring checks"`
	DefaultGroup    string `short:"g" long:"default-group" description:"default contact group to use if there is no valid group provided"`
	Config          string `short:"c" long:"config-file" default:"/etc/atlantis/supervisor/monitor.toml" description:"the config file to use"`
	TimeoutDuration uint   `short:"t" long:"timeout-duration" description:"max number of seconds to wait for a monitoring check to finish"`
	Verbose         bool   `short:"v" long:"verbose" default:false description:"print verbose debug information"`
}

type ServiceCheck struct {
	Service  string
	User     string
	Identity string
	Host     string
	Port     uint16
	Script   string
}

//TODO(mchandra):Need defaults defined by constants
var config = &Config{
	ContainerFile:   "/etc/atlantis/supervisor/save/containers",
	ContainersDir:   "/etc/atlantis/containers",
	InventoryDir:    "/etc/atlantis/supervisor/inventory",
	SSHIdentity:     "/opt/atlantis/supervisor/master_id_rsa",
	SSHUser:         "root",
	CheckName:       "ContainerMonitor",
	CheckDir:        "/check_mk_checks",
	DefaultGroup:    "atlantis_orphan_apps",
	TimeoutDuration: 110,
	Verbose:         false,
}

func (s *ServiceCheck) cmd() *exec.Cmd {
	return silentSshCmd(s.User, s.Identity, s.Host, s.Script, s.Port)
}

func (s *ServiceCheck) timeOutMsg() string {
	return fmt.Sprintf("%d %s - Timeout occured during check\n", Critical, s.Service)
}

func (s *ServiceCheck) errMsg(err error) string {
	if err != nil {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, err.Error())
	} else {
		return fmt.Sprintf("%d %s - Error encountered while monitoring the service\n", Critical, s.Service)
	}
}

func (s *ServiceCheck) validate(msg string) string {
	m := strings.SplitN(msg, " ", 4)
	if len(m) > 1 && m[1] == s.Service {
		return msg
	}
	return s.errMsg(nil)
}

func (s *ServiceCheck) runCheck(done chan bool) {
	out, err := s.cmd().Output()
	if err != nil {
		fmt.Print(s.errMsg(err))
	} else {
		fmt.Print(s.validate(string(out)))
	}
	done <- true
}

func (s *ServiceCheck) checkWithTimeout(results chan bool, d time.Duration) {
	done := make(chan bool, 1)
	go s.runCheck(done)
	select {
	case <-done:
		results <- true
	case <-time.After(d):
		fmt.Print(s.timeOutMsg())
		results <- true
	}
}

type ContainerCheck struct {
	Name         string
	User         string
	Identity     string
	Directory    string
	Inventory    string
	ContactGroup string
	container    *types.Container
}

type ContainerConfig struct {
	Dependencies map[string]interface{}
}

func (c *ContainerCheck) verifyContactGroup(group string) bool {
	output, err := exec.Command("/usr/bin/cmk_admin", "-l").Output()
	if err != nil {
		fmt.Printf("%d %s - Error listing existing contact_groups for validation, please try again later! Error: %s\n", Critical, config.CheckName, err.Error())
		return false
	}
	for _, l := range strings.Split(string(output), "\n") {
		cg := strings.TrimSpace(strings.TrimPrefix(l, "*"))
		if cg == c.ContactGroup {
			return true
		}
	}
	return false
}

func (c *ContainerCheck) parseContactGroup() {
	c.ContactGroup = config.DefaultGroup
	config_file := filepath.Join(config.ContainersDir, c.container.ID, "config.json")
	var cont_config ContainerConfig
	if err := serialize.RetrieveObject(config_file, &cont_config); err != nil {
		fmt.Printf("%d %s - Could not retrieve container config %s: %s\n", Critical, config.CheckName, config_file, err)
	} else {
		dep, ok := cont_config.Dependencies["cmk"]
		if !ok {
			fmt.Printf("%d %s - cmk dep not present, defaulting to %s contact group!\n", OK, config.CheckName, config.DefaultGroup)
			return
		}
		cmk_dep, ok := dep.(map[string]interface{})
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but value is not map[string]string!\n", Critical, config.CheckName)
			return
		}
		val, ok := cmk_dep["contact_group"]
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but no contact_group key!\n", Critical, config.CheckName)
			return
		}
		group, ok := val.(string)
		if ok {
			group = strings.ToLower(group)
			if c.verifyContactGroup(group) {
				c.ContactGroup = group
			} else {
				fmt.Printf("%d %s - Specified contact_group does not exist in cmk! Falling back to default group %s.\n", Critical, config.CheckName, config.DefaultGroup)
			}
		} else {
			fmt.Printf("%d %s - Value for contact_group key of cmk dep is not a string!\n", Critical, config.CheckName)
		}
	}
}

func (c *ContainerCheck) updateContactGroup(name string) {
	if len(c.ContactGroup) == 0 {
		c.parseContactGroup()
	}
	inventoryPath := path.Join(c.Inventory, name)
	if _, err := os.Stat(inventoryPath); os.IsNotExist(err) {
		output, err := exec.Command("/usr/bin/cmk_admin", "-s", name, "-a", c.ContactGroup).CombinedOutput()
		if err != nil {
			fmt.Printf("%d %s - Failure to update contact group for service %s. Error: %s\n", OK, config.CheckName, name, err.Error())
		} else {
			os.Create(inventoryPath)
		}
		if config.Verbose {
			fmt.Printf("\n/usr/bin/cmk_admin -s %s -a %s\n%s\n\n", name, c.ContactGroup, output)
		}
	}
}

func (c *ContainerCheck) Run(t time.Duration, done chan bool) {
	c.updateContactGroup(c.Name)
	defer func() { done <- true }()
	o, err := silentSshCmd(c.User, c.Identity, c.container.Host, "ls "+c.Directory, c.container.SSHPort).Output()
	if err != nil {
		fmt.Printf("%d %s - Error getting checks for container: %s\n", Critical, c.Name, err.Error())
		return
	}
	fmt.Printf("%d %s - Got checks for container\n", OK, c.Name)
	scripts := strings.Split(strings.TrimSpace(string(o)), "\n")
	if len(scripts) == 0 || len(scripts[0]) == 0 {
		// nothing to check on this container, exit
		return
	}
	c.checkAll(scripts, t)
}

func (c *ContainerCheck) checkAll(scripts []string, t time.Duration) {
	results := make(chan bool, len(scripts))
	for _, s := range scripts {
		serviceName := fmt.Sprintf("%s_%s", strings.Split(s, ".")[0], c.container.ID)
		c.updateContactGroup(serviceName)
		go c.serviceCheck(s).checkWithTimeout(results, t)
	}
	for _ = range scripts {
		<-results
	}
}

func (c *ContainerCheck) serviceCheck(script string) *ServiceCheck {
	// The full path to the script is required
	command := fmt.Sprintf("%s/%s %d %s", c.Directory, script, c.container.PrimaryPort, c.container.ID)
	// The service name is obtained be removing the file extension from the script and appending the container
	// id
	serviceName := fmt.Sprintf("%s_%s", strings.Split(script, ".")[0], c.container.ID)
	return &ServiceCheck{serviceName, c.User, c.Identity, c.container.Host, c.container.SSHPort, command}
}

func silentSshCmd(user, identity, host, cmd string, port uint16) *exec.Cmd {
	args := []string{"-q", user + "@" + host, "-i", identity, "-p", fmt.Sprintf("%d", port), "-o", "StrictHostKeyChecking=no", cmd}
	return exec.Command("ssh", args...)
}

func overlayConfig() {
	opts := &Opts{}
	flags.Parse(opts)
	if opts.Config != "" {
		_, err := toml.DecodeFile(opts.Config, config)
		if err != nil {
			// no need to panic here. we have reasonable defaults.
		}
	}
	if opts.ContainerFile != "" {
		config.ContainerFile = opts.ContainerFile
	}
	if opts.ContainersDir != "" {
		config.ContainersDir = opts.ContainersDir
	}
	if opts.SSHIdentity != "" {
		config.SSHIdentity = opts.SSHIdentity
	}
	if opts.SSHUser != "" {
		config.SSHUser = opts.SSHUser
	}
	if opts.CheckDir != "" {
		config.CheckDir = opts.CheckDir
	}
	if opts.CheckName != "" {
		config.CheckName = opts.CheckName
	}
	if opts.DefaultGroup != "" {
		config.DefaultGroup = opts.DefaultGroup
	}
	if opts.TimeoutDuration != 0 {
		config.TimeoutDuration = opts.TimeoutDuration
	}
	if opts.Verbose {
		config.Verbose = true
	}
}

//file containing containers and service name to show in Nagios for the monitor itself
func Run() {
	overlayConfig()
	var contMap map[string]*types.Container
	//Check if folder exists
	_, err := os.Stat(config.ContainerFile)
	if os.IsNotExist(err) {
		fmt.Printf("%d %s - Container file does not exists %s. Likely no live containers present.\n", OK, config.CheckName, config.ContainerFile)
		return
	}
	if err := serialize.RetrieveObject(config.ContainerFile, &contMap); err != nil {
		fmt.Printf("%d %s - Error retrieving %s: %s\n", Critical, config.CheckName, config.ContainerFile, err)
		return
	}
	done := make(chan bool, len(contMap))
	config.SSHIdentity = strings.Replace(config.SSHIdentity, "~", os.Getenv("HOME"), 1)
	for _, c := range contMap {
		if c.Host == "" {
			c.Host = "localhost"
		}
		check := &ContainerCheck{config.CheckName + "_" + c.ID, config.SSHUser, config.SSHIdentity, config.CheckDir, config.InventoryDir, "", c}
		go check.Run(time.Duration(config.TimeoutDuration)*time.Second, done)
	}
	for _ = range contMap {
		<-done
	}
	// Clean up inventories from containers that no longer exist
	err = filepath.Walk(config.InventoryDir, func(path string, _ os.FileInfo, _ error) error {
		if path == config.InventoryDir {
			return nil
		}
		var err error
		split := strings.Split(path, "_")
		cont := split[len(split)-1]
		if _, ok := contMap[cont]; !ok {
			err = os.Remove(path)
		}
		return err
	})
	if err != nil {
		fmt.Printf("%d %s - Error iterating over inventory to delete obsolete markers. Error: %s\n", OK, config.CheckName, err.Error())
	}
}
<commit_msg>have output alert correct service<commit_after>/* Copyright 2014 Ooyala, Inc. All rights reserved.
 *
 * This file is licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
 * except in compliance with the License. You may obtain a copy of the License at
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under the License is
 * distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and limitations under the License.
 */

package monitor

import (
	"atlantis/supervisor/containers/serialize"
	"atlantis/supervisor/rpc/types"
	"fmt"
	"github.com/BurntSushi/toml"
	"github.com/jigish/go-flags"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"
	"time"
)

const (
	OK = iota
	Warning
	Critical
	Uknown
)

type Config struct {
	ContainerFile   string `toml:"container_file"`
	ContainersDir   string `toml:"container_dir"`
	InventoryDir    string `toml:"inventory_dir"`
	SSHIdentity     string `toml:"ssh_identity"`
	SSHUser         string `toml:"ssh_user"`
	CheckName       string `toml:"check_name"`
	CheckDir        string `toml:"check_dir"`
	DefaultGroup    string `toml:"default_group"`
	TimeoutDuration uint   `toml:"timeout_duration"`
	Verbose         bool   `toml:"verbose"`
}

type Opts struct {
	ContainerFile   string `short:"f" long:"container-file" description:"file to get container information"`
	ContainersDir   string `short:"s" long:"containers-dir" description:"directory containing configs for each container"`
	SSHIdentity     string `short:"i" long:"ssh-identity" description:"file containing the SSH key for all containers"`
	SSHUser         string `short:"u" long:"ssh-user" description:"user account to ssh into containers"`
	CheckName       string `short:"n" long:"check-name" description:"service name that will appear in Nagios for the monitor"`
	CheckDir        string `short:"d" long:"check-dir" description:"directory containing all the scripts for the monitoring checks"`
	DefaultGroup    string `short:"g" long:"default-group" description:"default contact group to use if there is no valid group provided"`
	Config          string `short:"c" long:"config-file" default:"/etc/atlantis/supervisor/monitor.toml" description:"the config file to use"`
	TimeoutDuration uint   `short:"t" long:"timeout-duration" description:"max number of seconds to wait for a monitoring check to finish"`
	Verbose         bool   `short:"v" long:"verbose" default:false description:"print verbose debug information"`
}

type ServiceCheck struct {
	Service  string
	User     string
	Identity string
	Host     string
	Port     uint16
	Script   string
}

//TODO(mchandra):Need defaults defined by constants
var config = &Config{
	ContainerFile:   "/etc/atlantis/supervisor/save/containers",
	ContainersDir:   "/etc/atlantis/containers",
	InventoryDir:    "/etc/atlantis/supervisor/inventory",
	SSHIdentity:     "/opt/atlantis/supervisor/master_id_rsa",
	SSHUser:         "root",
	CheckName:       "ContainerMonitor",
	CheckDir:        "/check_mk_checks",
	DefaultGroup:    "atlantis_orphan_apps",
	TimeoutDuration: 110,
	Verbose:         false,
}

func (s *ServiceCheck) cmd() *exec.Cmd {
	return silentSshCmd(s.User, s.Identity, s.Host, s.Script, s.Port)
}

func (s *ServiceCheck) timeOutMsg() string {
	return fmt.Sprintf("%d %s - Timeout occured during check\n", Critical, s.Service)
}

func (s *ServiceCheck) errMsg(err error) string {
	if err != nil {
		return fmt.Sprintf("%d %s - %s\n", Critical, s.Service, err.Error())
	} else {
		return fmt.Sprintf("%d %s - Error encountered while monitoring the service\n", Critical, s.Service)
	}
}

func (s *ServiceCheck) validate(msg string) string {
	m := strings.SplitN(msg, " ", 4)
	if len(m) > 1 && m[1] == s.Service {
		return msg
	}
	return s.errMsg(nil)
}

func (s *ServiceCheck) runCheck(done chan bool) {
	out, err := s.cmd().Output()
	if err != nil {
		fmt.Print(s.errMsg(err))
	} else {
		fmt.Print(s.validate(string(out)))
	}
	done <- true
}

func (s *ServiceCheck) checkWithTimeout(results chan bool, d time.Duration) {
	done := make(chan bool, 1)
	go s.runCheck(done)
	select {
	case <-done:
		results <- true
	case <-time.After(d):
		fmt.Print(s.timeOutMsg())
		results <- true
	}
}

type ContainerCheck struct {
	Name         string
	User         string
	Identity     string
	Directory    string
	Inventory    string
	ContactGroup string
	container    *types.Container
}

type ContainerConfig struct {
	Dependencies map[string]interface{}
}

func (c *ContainerCheck) verifyContactGroup(group string) bool {
	output, err := exec.Command("/usr/bin/cmk_admin", "-l").Output()
	if err != nil {
		fmt.Printf("%d %s - Error listing existing contact_groups for validation, please try again later! Error: %s\n", Critical, c.Name, err.Error())
		return false
	}
	for _, l := range strings.Split(string(output), "\n") {
		cg := strings.TrimSpace(strings.TrimPrefix(l, "*"))
		if cg == c.ContactGroup {
			return true
		}
	}
	return false
}

func (c *ContainerCheck) parseContactGroup() {
	c.ContactGroup = config.DefaultGroup
	config_file := filepath.Join(config.ContainersDir, c.container.ID, "config.json")
	var cont_config ContainerConfig
	if err := serialize.RetrieveObject(config_file, &cont_config); err != nil {
		fmt.Printf("%d %s - Could not retrieve container config %s: %s\n", Critical, c.Name, config_file, err)
	} else {
		dep, ok := cont_config.Dependencies["cmk"]
		if !ok {
			fmt.Printf("%d %s - cmk dep not present, defaulting to %s contact group!\n", OK, c.Name, config.DefaultGroup)
			return
		}
		cmk_dep, ok := dep.(map[string]interface{})
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but value is not map[string]string!\n", Critical, c.Name)
			return
		}
		val, ok := cmk_dep["contact_group"]
		if !ok {
			fmt.Printf("%d %s - cmk dep present, but no contact_group key!\n", Critical, c.Name)
			return
		}
		group, ok := val.(string)
		if ok {
			group = strings.ToLower(group)
			if c.verifyContactGroup(group) {
				c.ContactGroup = group
			} else {
				fmt.Printf("%d %s - Specified contact_group does not exist in cmk! Falling back to default group %s.\n", Critical, c.Name, config.DefaultGroup)
			}
		} else {
			fmt.Printf("%d %s - Value for contact_group key of cmk dep is not a string!\n", Critical, c.Name)
		}
	}
}

func (c *ContainerCheck) updateContactGroup(name string) {
	if len(c.ContactGroup) == 0 {
		c.parseContactGroup()
	}
	inventoryPath := path.Join(c.Inventory, name)
	if _, err := os.Stat(inventoryPath); os.IsNotExist(err) {
		output, err := exec.Command("/usr/bin/cmk_admin", "-s", name, "-a", c.ContactGroup).CombinedOutput()
		if err != nil {
			fmt.Printf("%d %s - Failure to update contact group for service %s. Error: %s\n", OK, c.Name, name, err.Error())
		} else {
			os.Create(inventoryPath)
		}
		if config.Verbose {
			fmt.Printf("\n/usr/bin/cmk_admin -s %s -a %s\n%s\n\n", name, c.ContactGroup, output)
		}
	}
}

func (c *ContainerCheck) Run(t time.Duration, done chan bool) {
	c.updateContactGroup(c.Name)
	defer func() { done <- true }()
	o, err := silentSshCmd(c.User, c.Identity, c.container.Host, "ls "+c.Directory, c.container.SSHPort).Output()
	if err != nil {
		fmt.Printf("%d %s - Error getting checks for container: %s\n", Critical, c.Name, err.Error())
		return
	}
	fmt.Printf("%d %s - Got checks for container\n", OK, c.Name)
	scripts := strings.Split(strings.TrimSpace(string(o)), "\n")
	if len(scripts) == 0 || len(scripts[0]) == 0 {
		// nothing to check on this container, exit
		return
	}
	c.checkAll(scripts, t)
}

func (c *ContainerCheck) checkAll(scripts []string, t time.Duration) {
	results := make(chan bool, len(scripts))
	for _, s := range scripts {
		serviceName := fmt.Sprintf("%s_%s", strings.Split(s, ".")[0], c.container.ID)
		c.updateContactGroup(serviceName)
		go c.serviceCheck(s).checkWithTimeout(results, t)
	}
	for _ = range scripts {
		<-results
	}
}

func (c *ContainerCheck) serviceCheck(script string) *ServiceCheck {
	// The full path to the script is required
	command := fmt.Sprintf("%s/%s %d %s", c.Directory, script, c.container.PrimaryPort, c.container.ID)
	// The service name is obtained be removing the file extension from the script and appending the container
	// id
	serviceName := fmt.Sprintf("%s_%s", strings.Split(script, ".")[0], c.container.ID)
	return &ServiceCheck{serviceName, c.User, c.Identity, c.container.Host, c.container.SSHPort, command}
}

func silentSshCmd(user, identity, host, cmd string, port uint16) *exec.Cmd {
	args := []string{"-q", user + "@" + host, "-i", identity, "-p", fmt.Sprintf("%d", port), "-o", "StrictHostKeyChecking=no", cmd}
	return exec.Command("ssh", args...)
}

func overlayConfig() {
	opts := &Opts{}
	flags.Parse(opts)
	if opts.Config != "" {
		_, err := toml.DecodeFile(opts.Config, config)
		if err != nil {
			// no need to panic here. we have reasonable defaults.
		}
	}
	if opts.ContainerFile != "" {
		config.ContainerFile = opts.ContainerFile
	}
	if opts.ContainersDir != "" {
		config.ContainersDir = opts.ContainersDir
	}
	if opts.SSHIdentity != "" {
		config.SSHIdentity = opts.SSHIdentity
	}
	if opts.SSHUser != "" {
		config.SSHUser = opts.SSHUser
	}
	if opts.CheckDir != "" {
		config.CheckDir = opts.CheckDir
	}
	if opts.CheckName != "" {
		config.CheckName = opts.CheckName
	}
	if opts.DefaultGroup != "" {
		config.DefaultGroup = opts.DefaultGroup
	}
	if opts.TimeoutDuration != 0 {
		config.TimeoutDuration = opts.TimeoutDuration
	}
	if opts.Verbose {
		config.Verbose = true
	}
}

//file containing containers and service name to show in Nagios for the monitor itself
func Run() {
	overlayConfig()
	var contMap map[string]*types.Container
	//Check if folder exists
	_, err := os.Stat(config.ContainerFile)
	if os.IsNotExist(err) {
		fmt.Printf("%d %s - Container file does not exists %s. Likely no live containers present.\n", OK, config.CheckName, config.ContainerFile)
		return
	}
	if err := serialize.RetrieveObject(config.ContainerFile, &contMap); err != nil {
		fmt.Printf("%d %s - Error retrieving %s: %s\n", Critical, config.CheckName, config.ContainerFile, err)
		return
	}
	done := make(chan bool, len(contMap))
	config.SSHIdentity = strings.Replace(config.SSHIdentity, "~", os.Getenv("HOME"), 1)
	for _, c := range contMap {
		if c.Host == "" {
			c.Host = "localhost"
		}
		check := &ContainerCheck{config.CheckName + "_" + c.ID, config.SSHUser, config.SSHIdentity, config.CheckDir, config.InventoryDir, "", c}
		go check.Run(time.Duration(config.TimeoutDuration)*time.Second, done)
	}
	for _ = range contMap {
		<-done
	}
	// Clean up inventories from containers that no longer exist
	err = filepath.Walk(config.InventoryDir, func(path string, _ os.FileInfo, _ error) error {
		if path == config.InventoryDir {
			return nil
		}
		var err error
		split := strings.Split(path, "_")
		cont := split[len(split)-1]
		if _, ok := contMap[cont]; !ok {
			err = os.Remove(path)
		}
		return err
	})
	if err != nil {
		fmt.Printf("%d %s - Error iterating over inventory to delete obsolete markers. Error: %s\n", OK, config.CheckName, err.Error())
	}
}
<|endoftext|>


<commit_before>package sqlstore

import (
	"context"
	"testing"

	"github.com/grafana/grafana/pkg/components/simplejson"
	m "github.com/grafana/grafana/pkg/models"
	. "github.com/smartystreets/goconvey/convey"
)

func TestAlertNotificationSQLAccess(t *testing.T) {
	Convey("Testing Alert notification sql access", t, func() {
		InitTestDB(t)

		Convey("Alert notification journal", func() {
			var alertId int64 = 5
			var orgId int64 = 5
			var notifierId int64 = 5

			Convey("Getting last journal should raise error if no one exists", func() {
				query := &m.GetLatestNotificationQuery{AlertId: alertId, OrgId: orgId, NotifierId: notifierId}
				err := GetLatestNotification(context.Background(), query)
				So(err, ShouldEqual, m.ErrJournalingNotFound)

				Convey("shoulbe be able to record two journaling events", func() {
					createCmd := &m.RecordNotificationJournalCommand{AlertId: alertId, NotifierId: notifierId, OrgId: orgId, Success: true, SentAt: 1}

					err := RecordNotificationJournal(context.Background(), createCmd)
					So(err, ShouldBeNil)

					createCmd.SentAt += 1000 //increase epoch

					err = RecordNotificationJournal(context.Background(), createCmd)
					So(err, ShouldBeNil)

					Convey("get last journaling event", func() {
						err := GetLatestNotification(context.Background(), query)
						So(err, ShouldBeNil)
						So(query.Result.SentAt, ShouldEqual, 1001)

						Convey("be able to clear all journaling for an notifier", func() {
							cmd := &m.CleanNotificationJournalCommand{AlertId: alertId, NotifierId: notifierId, OrgId: orgId}
							err := CleanNotificationJournal(context.Background(), cmd)
							So(err, ShouldBeNil)

							Convey("querying for last junaling should raise error", func() {
								query := &m.GetLatestNotificationQuery{AlertId: alertId, OrgId: orgId, NotifierId: notifierId}
								err := GetLatestNotification(context.Background(), query)
								So(err, ShouldEqual, m.ErrJournalingNotFound)
							})
						})
					})
				})
			})
		})

		Convey("Alert notifications should be empty", func() {
			cmd := &m.GetAlertNotificationsQuery{
				OrgId: 2,
				Name:  "email",
			}

			err := GetAlertNotifications(cmd)
			So(err, ShouldBeNil)
			So(cmd.Result, ShouldBeNil)
		})

		Convey("Cannot save alert notifier with send reminder = true", func() {
			cmd := &m.CreateAlertNotificationCommand{
				Name:         "ops",
				Type:         "email",
				OrgId:        1,
				SendReminder: true,
				Settings:     simplejson.New(),
			}

			Convey("and missing frequency", func() {
				err := CreateAlertNotificationCommand(cmd)
				So(err, ShouldEqual, m.ErrNotificationFrequencyNotFound)
			})

			Convey("invalid frequency", func() {
				cmd.Frequency = "invalid duration"

				err := CreateAlertNotificationCommand(cmd)
				So(err.Error(), ShouldEqual, "time: invalid duration invalid duration")
			})
		})

		Convey("Cannot update alert notifier with notitfyonce = false", func() {
			cmd := &m.CreateAlertNotificationCommand{
				Name:         "ops update",
				Type:         "email",
				OrgId:        1,
				SendReminder: false,
				Settings:     simplejson.New(),
			}

			err := CreateAlertNotificationCommand(cmd)
			So(err, ShouldBeNil)

			updateCmd := &m.UpdateAlertNotificationCommand{
				Id:           cmd.Result.Id,
				SendReminder: true,
			}

			Convey("and missing frequency", func() {
				err := UpdateAlertNotification(updateCmd)
				So(err, ShouldEqual, m.ErrNotificationFrequencyNotFound)
			})

			Convey("invalid frequency", func() {
				updateCmd.Frequency = "invalid duration"

				err := UpdateAlertNotification(updateCmd)
				So(err, ShouldNotBeNil)
				So(err.Error(), ShouldEqual, "time: invalid duration invalid duration")
			})
		})

		Convey("Can save Alert Notification", func() {
			cmd := &m.CreateAlertNotificationCommand{
				Name:         "ops",
				Type:         "email",
				OrgId:        1,
				SendReminder: true,
				Frequency:    "10s",
				Settings:     simplejson.New(),
			}

			err := CreateAlertNotificationCommand(cmd)
			So(err, ShouldBeNil)
			So(cmd.Result.Id, ShouldNotEqual, 0)
			So(cmd.Result.OrgId, ShouldNotEqual, 0)
			So(cmd.Result.Type, ShouldEqual, "email")

			Convey("Cannot save Alert Notification with the same name", func() {
				err = CreateAlertNotificationCommand(cmd)
				So(err, ShouldNotBeNil)
			})

			Convey("Can update alert notification", func() {
				newCmd := &m.UpdateAlertNotificationCommand{
					Name:         "NewName",
					Type:         "webhook",
					OrgId:        cmd.Result.OrgId,
					SendReminder: true,
					Frequency:    "10s",
					Settings:     simplejson.New(),
					Id:           cmd.Result.Id,
				}
				err := UpdateAlertNotification(newCmd)
				So(err, ShouldBeNil)
				So(newCmd.Result.Name, ShouldEqual, "NewName")
			})
		})

		Convey("Can search using an array of ids", func() {
			cmd1 := m.CreateAlertNotificationCommand{Name: "nagios", Type: "webhook", OrgId: 1, SendReminder: true, Frequency: "10s", Settings: simplejson.New()}
			cmd2 := m.CreateAlertNotificationCommand{Name: "slack", Type: "webhook", OrgId: 1, SendReminder: true, Frequency: "10s", Settings: simplejson.New()}
			cmd3 := m.CreateAlertNotificationCommand{Name: "ops2", Type: "email", OrgId: 1, SendReminder: true, Frequency: "10s", Settings: simplejson.New()}
			cmd4 := m.CreateAlertNotificationCommand{IsDefault: true, Name: "default", Type: "email", OrgId: 1, SendReminder: true, Frequency: "10s", Settings: simplejson.New()}

			otherOrg := m.CreateAlertNotificationCommand{Name: "default", Type: "email", OrgId: 2, SendReminder: true, Frequency: "10s", Settings: simplejson.New()}

			So(CreateAlertNotificationCommand(&cmd1), ShouldBeNil)
			So(CreateAlertNotificationCommand(&cmd2), ShouldBeNil)
			So(CreateAlertNotificationCommand(&cmd3), ShouldBeNil)
			So(CreateAlertNotificationCommand(&cmd4), ShouldBeNil)
			So(CreateAlertNotificationCommand(&otherOrg), ShouldBeNil)

			Convey("search", func() {
				query := &m.GetAlertNotificationsToSendQuery{
					Ids:   []int64{cmd1.Result.Id, cmd2.Result.Id, 112341231},
					OrgId: 1,
				}

				err := GetAlertNotificationsToSend(query)
				So(err, ShouldBeNil)
				So(len(query.Result), ShouldEqual, 3)
			})

			Convey("all", func() {
				query := &m.GetAllAlertNotificationsQuery{
					OrgId: 1,
				}

				err := GetAllAlertNotifications(query)
				So(err, ShouldBeNil)
				So(len(query.Result), ShouldEqual, 4)
			})
		})
	})
}
<commit_msg>change/add tests for alerting notification reminders<commit_after>package sqlstore

import (
	"context"
	"testing"
	"time"

	"github.com/grafana/grafana/pkg/components/simplejson"
	m "github.com/grafana/grafana/pkg/models"
	. "github.com/smartystreets/goconvey/convey"
)

func TestAlertNotificationSQLAccess(t *testing.T) {
	Convey("Testing Alert notification sql access", t, func() {
		InitTestDB(t)

		Convey("Alert notification journal", func() {
			var alertId int64 = 5
			var orgId int64 = 5
			var notifierId int64 = 5

			Convey("Getting last journal should raise error if no one exists", func() {
				query := &m.GetLatestNotificationQuery{AlertId: alertId, OrgId: orgId, NotifierId: notifierId}
				err := GetLatestNotification(context.Background(), query)
				So(err, ShouldEqual, m.ErrJournalingNotFound)

				Convey("shoulbe be able to record two journaling events", func() {
					createCmd := &m.RecordNotificationJournalCommand{AlertId: alertId, NotifierId: notifierId, OrgId: orgId, Success: true, SentAt: 1}

					err := RecordNotificationJournal(context.Background(), createCmd)
					So(err, ShouldBeNil)

					createCmd.SentAt += 1000 //increase epoch

					err = RecordNotificationJournal(context.Background(), createCmd)
					So(err, ShouldBeNil)

					Convey("get last journaling event", func() {
						err := GetLatestNotification(context.Background(), query)
						So(err, ShouldBeNil)
						So(query.Result.SentAt, ShouldEqual, 1001)

						Convey("be able to clear all journaling for an notifier", func() {
							cmd := &m.CleanNotificationJournalCommand{AlertId: alertId, NotifierId: notifierId, OrgId: orgId}
							err := CleanNotificationJournal(context.Background(), cmd)
							So(err, ShouldBeNil)

							Convey("querying for last junaling should raise error", func() {
								query := &m.GetLatestNotificationQuery{AlertId: alertId, OrgId: orgId, NotifierId: notifierId}
								err := GetLatestNotification(context.Background(), query)
								So(err, ShouldEqual, m.ErrJournalingNotFound)
							})
						})
					})
				})
			})
		})

		Convey("Alert notifications should be empty", func() {
			cmd := &m.GetAlertNotificationsQuery{
				OrgId: 2,
				Name:  "email",
			}

			err := GetAlertNotifications(cmd)
			So(err, ShouldBeNil)
			So(cmd.Result, ShouldBeNil)
		})

		Convey("Cannot save alert notifier with send reminder = true", func() {
			cmd := &m.CreateAlertNotificationCommand{
				Name:         "ops",
				Type:         "email",
				OrgId:        1,
				SendReminder: true,
				Settings:     simplejson.New(),
			}

			Convey("and missing frequency", func() {
				err := CreateAlertNotificationCommand(cmd)
				So(err, ShouldEqual, m.ErrNotificationFrequencyNotFound)
			})

			Convey("invalid frequency", func() {
				cmd.Frequency = "invalid duration"

				err := CreateAlertNotificationCommand(cmd)
				So(err.Error(), ShouldEqual, "time: invalid duration invalid duration")
			})
		})

		Convey("Cannot update alert notifier with send reminder = false", func() {
			cmd := &m.CreateAlertNotificationCommand{
				Name:         "ops update",
				Type:         "email",
				OrgId:        1,
				SendReminder: false,
				Settings:     simplejson.New(),
			}

			err := CreateAlertNotificationCommand(cmd)
			So(err, ShouldBeNil)

			updateCmd := &m.UpdateAlertNotificationCommand{
				Id:           cmd.Result.Id,
				SendReminder: true,
			}

			Convey("and missing frequency", func() {
				err := UpdateAlertNotification(updateCmd)
				So(err, ShouldEqual, m.ErrNotificationFrequencyNotFound)
			})

			Convey("invalid frequency", func() {
				updateCmd.Frequency = "invalid duration"

				err := UpdateAlertNotification(updateCmd)
				So(err, ShouldNotBeNil)
				So(err.Error(), ShouldEqual, "time: invalid duration invalid duration")
			})
		})

		Convey("Can save Alert Notification", func() {
			cmd := &m.CreateAlertNotificationCommand{
				Name:         "ops",
				Type:         "email",
				OrgId:        1,
				SendReminder: true,
				Frequency:    "10s",
				Settings:     simplejson.New(),
			}

			err := CreateAlertNotificationCommand(cmd)
			So(err, ShouldBeNil)
			So(cmd.Result.Id, ShouldNotEqual, 0)
			So(cmd.Result.OrgId, ShouldNotEqual, 0)
			So(cmd.Result.Type, ShouldEqual, "email")
			So(cmd.Result.Frequency, ShouldEqual, 10*time.Second)

			Convey("Cannot save Alert Notification with the same name", func() {
				err = CreateAlertNotificationCommand(cmd)
				So(err, ShouldNotBeNil)
			})

			Convey("Can update alert notification", func() {
				newCmd := &m.UpdateAlertNotificationCommand{
					Name:         "NewName",
					Type:         "webhook",
					OrgId:        cmd.Result.OrgId,
					SendReminder: true,
					Frequency:    "60s",
					Settings:     simplejson.New(),
					Id:           cmd.Result.Id,
				}
				err := UpdateAlertNotification(newCmd)
				So(err, ShouldBeNil)
				So(newCmd.Result.Name, ShouldEqual, "NewName")
				So(newCmd.Result.Frequency, ShouldEqual, 60*time.Second)
			})

			Convey("Can update alert notification to disable sending of reminders", func() {
				newCmd := &m.UpdateAlertNotificationCommand{
					Name:         "NewName",
					Type:         "webhook",
					OrgId:        cmd.Result.OrgId,
					SendReminder: false,
					Settings:     simplejson.New(),
					Id:           cmd.Result.Id,
				}
				err := UpdateAlertNotification(newCmd)
				So(err, ShouldBeNil)
				So(newCmd.Result.SendReminder, ShouldBeFalse)
			})
		})

		Convey("Can search using an array of ids", func() {
			cmd1 := m.CreateAlertNotificationCommand{Name: "nagios", Type: "webhook", OrgId: 1, SendReminder: true, Frequency: "10s", Settings: simplejson.New()}
			cmd2 := m.CreateAlertNotificationCommand{Name: "slack", Type: "webhook", OrgId: 1, SendReminder: true, Frequency: "10s", Settings: simplejson.New()}
			cmd3 := m.CreateAlertNotificationCommand{Name: "ops2", Type: "email", OrgId: 1, SendReminder: true, Frequency: "10s", Settings: simplejson.New()}
			cmd4 := m.CreateAlertNotificationCommand{IsDefault: true, Name: "default", Type: "email", OrgId: 1, SendReminder: true, Frequency: "10s", Settings: simplejson.New()}

			otherOrg := m.CreateAlertNotificationCommand{Name: "default", Type: "email", OrgId: 2, SendReminder: true, Frequency: "10s", Settings: simplejson.New()}

			So(CreateAlertNotificationCommand(&cmd1), ShouldBeNil)
			So(CreateAlertNotificationCommand(&cmd2), ShouldBeNil)
			So(CreateAlertNotificationCommand(&cmd3), ShouldBeNil)
			So(CreateAlertNotificationCommand(&cmd4), ShouldBeNil)
			So(CreateAlertNotificationCommand(&otherOrg), ShouldBeNil)

			Convey("search", func() {
				query := &m.GetAlertNotificationsToSendQuery{
					Ids:   []int64{cmd1.Result.Id, cmd2.Result.Id, 112341231},
					OrgId: 1,
				}

				err := GetAlertNotificationsToSend(query)
				So(err, ShouldBeNil)
				So(len(query.Result), ShouldEqual, 3)
			})

			Convey("all", func() {
				query := &m.GetAllAlertNotificationsQuery{
					OrgId: 1,
				}

				err := GetAllAlertNotifications(query)
				So(err, ShouldBeNil)
				So(len(query.Result), ShouldEqual, 4)
			})
		})
	})
}
<|endoftext|>


<commit_before>//Package svg implements the SVG backend for libcairo rendering.
package svg

//#cgo pkg-config: cairo
//#include <cairo/cairo.h>
//#include <cairo/cairo-svg.h>
import "C"

import (
	"io"
	"unsafe"

	"github.com/jimmyfrasche/cairo"
)

//Surface is an SVG surface.
//
//Surface implements cairo.VectorSurface.
//
//Originally cairo_svg_surface_t.
type Surface struct {
	cairo.ExtensionVectorSurface
	//w is used in NewWriter to ensure a reference to the writer lives as long as we do
	w io.Writer
}

//NewFile creates a new SVG surface that writes to filename.
//
//The parameters width and height are in the unit of a typographical point
//(1 point = 1/72 inch).
//
//Originally cairo_svg_surface_create.
func NewFile(filename string, width, height float64) (Surface, error) {
	nm := C.CString(filename)
	svg := C.cairo_svg_surface_create(nm, C.double(width), C.double(height))
	s := Surface{
		ExtensionVectorSurface: cairo.ExtensionNewVectorSurface(svg),
	}
	return s, s.Err()
}

//BUG(jmf): filename's CString not freed. Does cairo take control of it or strdup it?

//New creates a new SVG surface that writes to writer.
//
//If writer needs to be flushed or closed, that is the responsibility
//of the caller.
//
//The parameters width and height are in the unit of a typographical point
//(1 point = 1/72 inch).
//
//Originally cairo_svg_surface_create_for_stream.
func New(w io.Writer, width, height float64) (Surface, error) {
	wp := unsafe.Pointer(&w)
	svg := C.cairo_svg_surface_create_for_stream(cairo.ExtensionCairoWriteFuncT, wp, C.double(width), C.double(height))
	s := Surface{
		ExtensionVectorSurface: cairo.ExtensionNewVectorSurface(svg),
		w: w,
	}
	return s, s.Err()
}

func cNew(s *C.cairo_surface_t) (cairo.Surface, error) {
	//Note that if the surface was created with an io.Writer we have no way of
	//getting it here but that's okay as long as the original reference lives on.
	S := Surface{
		ExtensionVectorSurface: cairo.ExtensionNewVectorSurface(s),
	}
	return S, S.Err()
}

func init() {
	cairo.ExtensionRegisterRawToSurface(cairo.SurfaceTypeSVG, cNew)
}

//BUG(jmf): add documentation about mime type to New after I figure out what that entails.

//BUG(jmf): add special method for attaching url's as that's handled specially anyhoo.

//RestrictTo restricts the generated SVG file to the specified version.
//
//Originally cairo_svg_surface_restrict_to_version.
func (s Surface) RestrictTo(v version) error {
	C.cairo_svg_surface_restrict_to_version(s.ExtensionRaw(), v.c())
	return s.Err()
}
<commit_msg>fixed memory leak and updated comments<commit_after>//Package svg implements the SVG backend for libcairo rendering.
package svg

//#cgo pkg-config: cairo
//#include <stdlib.h>
//#include <cairo/cairo.h>
//#include <cairo/cairo-svg.h>
import "C"

import (
	"io"
	"unsafe"

	"github.com/jimmyfrasche/cairo"
)

//Surface is an SVG surface.
//
//Surface implements cairo.VectorSurface.
type Surface struct {
	cairo.ExtensionVectorSurface
	//w is used in NewWriter to ensure a reference to the writer lives as long as we do
	w io.Writer
}

//NewFile creates a new SVG surface that writes to filename.
//
//The parameters width and height are in the unit of a typographical point
//(1 point = 1/72 inch).
//
//Originally cairo_svg_surface_create.
func NewFile(filename string, width, height float64) (Surface, error) {
	nm := C.CString(filename)
	svg := C.cairo_svg_surface_create(nm, C.double(width), C.double(height))
	C.free(unsafe.Pointer(nm))
	s := Surface{
		ExtensionVectorSurface: cairo.ExtensionNewVectorSurface(svg),
	}
	return s, s.Err()
}

//New creates a new SVG surface that writes to writer.
//
//If writer needs to be flushed or closed, that is the responsibility
//of the caller.
//
//The parameters width and height are in the unit of a typographical point
//(1 point = 1/72 inch).
//
//Originally cairo_svg_surface_create_for_stream.
func New(w io.Writer, width, height float64) (Surface, error) {
	wp := unsafe.Pointer(&w)
	svg := C.cairo_svg_surface_create_for_stream(cairo.ExtensionCairoWriteFuncT, wp, C.double(width), C.double(height))
	s := Surface{
		ExtensionVectorSurface: cairo.ExtensionNewVectorSurface(svg),
		w: w,
	}
	return s, s.Err()
}

func cNew(s *C.cairo_surface_t) (cairo.Surface, error) {
	//Note that if the surface was created with an io.Writer we have no way of
	//getting it here but that's okay as long as the original reference lives on.
	S := Surface{
		ExtensionVectorSurface: cairo.ExtensionNewVectorSurface(s),
	}
	return S, S.Err()
}

func init() {
	cairo.ExtensionRegisterRawToSurface(cairo.SurfaceTypeSVG, cNew)
}

//BUG(jmf): add documentation about mime type to New after I figure out what that entails.

//BUG(jmf): add special method for attaching url's as that's handled specially anyhoo.

//RestrictTo restricts the generated SVG file to the specified version.
//
//This method should only be called before any drawing operations have been
//performed on this surface.
//
//Originally cairo_svg_surface_restrict_to_version.
func (s Surface) RestrictTo(v version) error {
	C.cairo_svg_surface_restrict_to_version(s.ExtensionRaw(), v.c())
	return s.Err()
}
<|endoftext|>


<commit_before>package svg // import "github.com/tdewolff/minify/svg"

import (
	"io"

	"github.com/tdewolff/buffer"
	"github.com/tdewolff/minify"
	xmlMinify "github.com/tdewolff/minify/xml"
	"github.com/tdewolff/parse"
	cssParse "github.com/tdewolff/parse/css"
	"github.com/tdewolff/parse/svg"
	"github.com/tdewolff/parse/xml"
)

var (
	ltBytes         = []byte("<")
	gtBytes         = []byte(">")
	voidBytes       = []byte("/>")
	isBytes         = []byte("=")
	spaceBytes      = []byte(" ")
	emptyBytes      = []byte("\"\"")
	endBytes        = []byte("</")
	CDATAStartBytes = []byte("<![CDATA[")
	CDATAEndBytes   = []byte("]]>")
)

////////////////////////////////////////////////////////////////

// Minify minifies XML files, it reads from r and writes to w.
// Removes unnecessary whitespace, tags, attributes, quotes and comments and typically saves 10% in size.
func Minify(m minify.Minifier, _ string, w io.Writer, r io.Reader) error {
	var tag svg.Hash

	attrMinifyBuffer := buffer.NewWriter(make([]byte, 0, 64))
	attrByteBuffer := make([]byte, 0, 64)

	z := xml.NewTokenizer(r)
	tb := xmlMinify.NewTokenBuffer(z)
	for {
		t := *tb.Shift()
		if t.TokenType == xml.CDATAToken {
			var useCDATA bool
			if t.Data, useCDATA = xmlMinify.EscapeCDATAVal(&attrByteBuffer, t.Data); !useCDATA {
				t.TokenType = xml.TextToken
			}
		}
		switch t.TokenType {
		case xml.ErrorToken:
			if z.Err() == io.EOF {
				return nil
			}
			return z.Err()
		case xml.TextToken:
			t.Data = parse.ReplaceMultiple(parse.Trim(t.Data, parse.IsWhitespace), parse.IsWhitespace, ' ')
			if tag == svg.Style && len(t.Data) > 0 {
				if err := m.Minify("text/css", w, buffer.NewReader(t.Data)); err != nil {
					if err == minify.ErrNotExist { // no minifier, write the original
						if _, err := w.Write(t.Data); err != nil {
							return err
						}
					} else {
						return err
					}
				}
			} else if _, err := w.Write(t.Data); err != nil {
				return err
			}
		case xml.CDATAToken:
			if _, err := w.Write(CDATAStartBytes); err != nil {
				return err
			}
			t.Data = parse.ReplaceMultiple(parse.Trim(t.Data, parse.IsWhitespace), parse.IsWhitespace, ' ')
			if tag == svg.Style && len(t.Data) > 0 {
				if err := m.Minify("text/css", w, buffer.NewReader(t.Data)); err != nil {
					if err == minify.ErrNotExist { // no minifier, write the original
						if _, err := w.Write(t.Data); err != nil {
							return err
						}
					} else {
						return err
					}
				}
			} else if _, err := w.Write(t.Data); err != nil {
				return err
			}
			if _, err := w.Write(CDATAEndBytes); err != nil {
				return err
			}
		case xml.StartTagToken:
			if _, err := w.Write(ltBytes); err != nil {
				return err
			}
			if _, err := w.Write(t.Data); err != nil {
				return err
			}
			tag = svg.ToHash(t.Data)
		case xml.StartTagPIToken:
			for {
				if t := *tb.Shift(); t.TokenType != xml.StartTagClosePIToken {
					break
				}
			}
		case xml.AttributeToken:
			if len(t.AttrVal) < 2 {
				continue
			}
			attr := svg.ToHash(t.Data)
			val := parse.ReplaceMultiple(parse.Trim(t.AttrVal[1:len(t.AttrVal)-1], parse.IsWhitespace), parse.IsWhitespace, ' ')
			if tag == svg.Svg && attr == svg.Version {
				continue
			}

			if _, err := w.Write(spaceBytes); err != nil {
				return err
			}
			if _, err := w.Write(t.Data); err != nil {
				return err
			}
			if _, err := w.Write(isBytes); err != nil {
				return err
			}

			if attr == svg.Style {
				attrMinifyBuffer.Reset()
				if m.Minify("text/css;inline=1", attrMinifyBuffer, buffer.NewReader(val)) == nil {
					val = attrMinifyBuffer.Bytes()
				}
			} else if tag == svg.Path && attr == svg.D {
				val = shortenPathData(val)
			} else if num, dim, ok := cssParse.SplitNumberDimension(val); ok {
				num = minify.MinifyNumber(num)
				if len(num) == 1 && num[0] == '0' {
					val = num
				} else {
					if len(dim) > 1 { // only percentage is length 1
						parse.ToLower(dim)
					}
					val = append(num, dim...)
				}
			}

			// prefer single or double quotes depending on what occurs more often in value
			val = xmlMinify.EscapeAttrVal(&attrByteBuffer, val)
			if _, err := w.Write(val); err != nil {
				return err
			}
		case xml.StartTagCloseToken:
			next := tb.Peek(0)
			skipExtra := false
			if next.TokenType == xml.TextToken && parse.IsAllWhitespace(next.Data) {
				next = tb.Peek(1)
				skipExtra = true
			}
			if next.TokenType == xml.EndTagToken {
				// collapse empty tags to single void tag
				tb.Shift()
				if skipExtra {
					tb.Shift()
				}
				if _, err := w.Write(voidBytes); err != nil {
					return err
				}
			} else {
				if _, err := w.Write(gtBytes); err != nil {
					return err
				}
			}
		case xml.StartTagCloseVoidToken:
			if _, err := w.Write(voidBytes); err != nil {
				return err
			}
		case xml.EndTagToken:
			if _, err := w.Write(endBytes); err != nil {
				return err
			}
			if _, err := w.Write(t.Data); err != nil {
				return err
			}
			if _, err := w.Write(gtBytes); err != nil {
				return err
			}
		}
	}
}

func shortenPathData(b []byte) []byte {
	cmd := byte(0)
	prevDigit := false
	prevDigitHasDecimal := false
	j := 0
	start := 0
	for i := 0; i < len(b); i++ {
		c := b[i]
		if c >= '0' && c <= '9' || c == '.' || c == '-' || c == '+' {
			hasDecimal := (c == '.')
			if start != 0 {
				j += copy(b[j:], b[start:i])
			} else {
				j += i
			}
			start = i
			i++
			for {
				c := b[i]
				if c == '.' && !hasDecimal {
					hasDecimal = true
				} else if c < '0' || c > '9' {
					break
				}
				i++
			}
			if c := b[i]; c == 'e' || c == 'E' {
				i++
				if c := b[i]; c == '+' || c == '-' {
					i++
				}
				for {
					if c := b[i]; c < '0' || c > '9' {
						break
					}
					i++
				}
			}
			num := minify.MinifyNumber(b[start:i])
			if prevDigit && (num[0] >= '0' && num[0] <= '9' || num[0] == '.' && !prevDigitHasDecimal) {
				b[j] = ' '
				j++
			}
			prevDigit = true
			prevDigitHasDecimal = false
			for _, c2 := range num {
				if c2 == '.' {
					prevDigitHasDecimal = true
					break
				}
			}
			j += copy(b[j:], num)
			start = i
			i--
		} else if c == ' ' || c == ',' || c == '\t' || c == '\n' || c == '\r' {
			if start != 0 {
				j += copy(b[j:], b[start:i])
			} else {
				j += i
			}
			start = i + 1
		} else {
			if cmd == c {
				if start != 0 {
					j += copy(b[j:], b[start:i])
				} else {
					j += i
				}
				start = i + 1
			} else {
				cmd = c
				prevDigit = false
			}
		}
	}
	if start != 0 {
		j += copy(b[j:], b[start:])
		return b[:j]
	}
	return b
}
<commit_msg>Revised path data minification<commit_after>package svg // import "github.com/tdewolff/minify/svg"

import (
	"io"

	"github.com/tdewolff/buffer"
	"github.com/tdewolff/minify"
	xmlMinify "github.com/tdewolff/minify/xml"
	"github.com/tdewolff/parse"
	cssParse "github.com/tdewolff/parse/css"
	"github.com/tdewolff/parse/svg"
	"github.com/tdewolff/parse/xml"
)

var (
	ltBytes         = []byte("<")
	gtBytes         = []byte(">")
	voidBytes       = []byte("/>")
	isBytes         = []byte("=")
	spaceBytes      = []byte(" ")
	emptyBytes      = []byte("\"\"")
	endBytes        = []byte("</")
	CDATAStartBytes = []byte("<![CDATA[")
	CDATAEndBytes   = []byte("]]>")
)

////////////////////////////////////////////////////////////////

// Minify minifies XML files, it reads from r and writes to w.
// Removes unnecessary whitespace, tags, attributes, quotes and comments and typically saves 10% in size.
func Minify(m minify.Minifier, _ string, w io.Writer, r io.Reader) error {
	var tag svg.Hash

	attrMinifyBuffer := buffer.NewWriter(make([]byte, 0, 64))
	attrByteBuffer := make([]byte, 0, 64)

	z := xml.NewTokenizer(r)
	tb := xmlMinify.NewTokenBuffer(z)
	for {
		t := *tb.Shift()
		if t.TokenType == xml.CDATAToken {
			var useCDATA bool
			if t.Data, useCDATA = xmlMinify.EscapeCDATAVal(&attrByteBuffer, t.Data); !useCDATA {
				t.TokenType = xml.TextToken
			}
		}
		switch t.TokenType {
		case xml.ErrorToken:
			if z.Err() == io.EOF {
				return nil
			}
			return z.Err()
		case xml.TextToken:
			t.Data = parse.ReplaceMultiple(parse.Trim(t.Data, parse.IsWhitespace), parse.IsWhitespace, ' ')
			if tag == svg.Style && len(t.Data) > 0 {
				if err := m.Minify("text/css", w, buffer.NewReader(t.Data)); err != nil {
					if err == minify.ErrNotExist { // no minifier, write the original
						if _, err := w.Write(t.Data); err != nil {
							return err
						}
					} else {
						return err
					}
				}
			} else if _, err := w.Write(t.Data); err != nil {
				return err
			}
		case xml.CDATAToken:
			if _, err := w.Write(CDATAStartBytes); err != nil {
				return err
			}
			t.Data = parse.ReplaceMultiple(parse.Trim(t.Data, parse.IsWhitespace), parse.IsWhitespace, ' ')
			if tag == svg.Style && len(t.Data) > 0 {
				if err := m.Minify("text/css", w, buffer.NewReader(t.Data)); err != nil {
					if err == minify.ErrNotExist { // no minifier, write the original
						if _, err := w.Write(t.Data); err != nil {
							return err
						}
					} else {
						return err
					}
				}
			} else if _, err := w.Write(t.Data); err != nil {
				return err
			}
			if _, err := w.Write(CDATAEndBytes); err != nil {
				return err
			}
		case xml.StartTagToken:
			if _, err := w.Write(ltBytes); err != nil {
				return err
			}
			if _, err := w.Write(t.Data); err != nil {
				return err
			}
			tag = svg.ToHash(t.Data)
		case xml.StartTagPIToken:
			for {
				if t := *tb.Shift(); t.TokenType != xml.StartTagClosePIToken {
					break
				}
			}
		case xml.AttributeToken:
			if len(t.AttrVal) < 2 {
				continue
			}
			attr := svg.ToHash(t.Data)
			val := parse.ReplaceMultiple(parse.Trim(t.AttrVal[1:len(t.AttrVal)-1], parse.IsWhitespace), parse.IsWhitespace, ' ')
			if tag == svg.Svg && attr == svg.Version {
				continue
			}

			if _, err := w.Write(spaceBytes); err != nil {
				return err
			}
			if _, err := w.Write(t.Data); err != nil {
				return err
			}
			if _, err := w.Write(isBytes); err != nil {
				return err
			}

			if attr == svg.Style {
				attrMinifyBuffer.Reset()
				if m.Minify("text/css;inline=1", attrMinifyBuffer, buffer.NewReader(val)) == nil {
					val = attrMinifyBuffer.Bytes()
				}
			} else if tag == svg.Path && attr == svg.D {
				val = shortenPathData(val)
			} else if num, dim, ok := cssParse.SplitNumberDimension(val); ok {
				num = minify.MinifyNumber(num)
				if len(num) == 1 && num[0] == '0' {
					val = num
				} else {
					if len(dim) > 1 { // only percentage is length 1
						parse.ToLower(dim)
					}
					val = append(num, dim...)
				}
			}

			// prefer single or double quotes depending on what occurs more often in value
			val = xmlMinify.EscapeAttrVal(&attrByteBuffer, val)
			if _, err := w.Write(val); err != nil {
				return err
			}
		case xml.StartTagCloseToken:
			next := tb.Peek(0)
			skipExtra := false
			if next.TokenType == xml.TextToken && parse.IsAllWhitespace(next.Data) {
				next = tb.Peek(1)
				skipExtra = true
			}
			if next.TokenType == xml.EndTagToken {
				// collapse empty tags to single void tag
				tb.Shift()
				if skipExtra {
					tb.Shift()
				}
				if _, err := w.Write(voidBytes); err != nil {
					return err
				}
			} else {
				if _, err := w.Write(gtBytes); err != nil {
					return err
				}
			}
		case xml.StartTagCloseVoidToken:
			if _, err := w.Write(voidBytes); err != nil {
				return err
			}
		case xml.EndTagToken:
			if _, err := w.Write(endBytes); err != nil {
				return err
			}
			if _, err := w.Write(t.Data); err != nil {
				return err
			}
			if _, err := w.Write(gtBytes); err != nil {
				return err
			}
		}
	}
}

func shortenPathData(b []byte) []byte {
	cmd := byte(0)
	prevDigit := false
	prevDigitRequiresSpace := true
	j := 0
	start := 0
	for i := 0; i < len(b); i++ {
		c := b[i]
		if c == ' ' || c == ',' || c == '\t' || c == '\n' || c == '\r' {
			if start != 0 {
				j += copy(b[j:], b[start:i])
			} else {
				j += i
			}
			start = i + 1
		} else if n, ok := cssParse.ParseNumber(b[i:]); ok {
			if start != 0 {
				j += copy(b[j:], b[start:i])
			} else {
				j += i
			}
			num := minify.MinifyNumber(b[i : i+n])
			if prevDigit && (num[0] >= '0' && num[0] <= '9' || num[0] == '.' && prevDigitRequiresSpace) {
				b[j] = ' '
				j++
			}
			prevDigit = true
			prevDigitRequiresSpace = true
			for _, c := range num {
				if c == '.' || c == 'e' || c == 'E' {
					prevDigitRequiresSpace = false
					break
				}
			}
			j += copy(b[j:], num)
			start = i + n
			i += n - 1
		} else {
			if cmd == c {
				if start != 0 {
					j += copy(b[j:], b[start:i])
				} else {
					j += i
				}
				start = i + 1
			} else {
				cmd = c
				prevDigit = false
			}
		}
	}
	if start != 0 {
		j += copy(b[j:], b[start:])
		return b[:j]
	}
	return b
}
<|endoftext|>


<commit_before>package mbr

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"fmt"
	"io"
	"io/ioutil"
	"strings"
	"testing"

	"github.com/diskfs/go-diskfs/testhelper"
)

const (
	mbrPartitionFile = "./testdata/mbr_partition.dat"
	partitionStart   = 2048
	partitionSize    = 20480
)

func TestFromBytes(t *testing.T) {
	t.Run("Short byte slice", func(t *testing.T) {
		b := make([]byte, partitionEntrySize-1, partitionEntrySize-1)
		rand.Read(b)
		partition, err := partitionFromBytes(b)
		if partition != nil {
			t.Error("should return nil partition")
		}
		if err == nil {
			t.Error("should not return nil error")
		}
		expected := fmt.Sprintf("Data for partition was %d bytes instead of expected %d", len(b), partitionEntrySize)
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("Long byte slice", func(t *testing.T) {
		b := make([]byte, partitionEntrySize+1, partitionEntrySize+1)
		rand.Read(b)
		partition, err := partitionFromBytes(b)
		if partition != nil {
			t.Error("should return nil partition")
		}
		if err == nil {
			t.Error("should not return nil error")
		}
		expected := fmt.Sprintf("Data for partition was %d bytes instead of expected %d", len(b), partitionEntrySize)
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("Invalid partition bootable code", func(t *testing.T) {
		b := make([]byte, partitionEntrySize, partitionEntrySize)
		rand.Read(b)
		b[0] = 0x67
		partition, err := partitionFromBytes(b)
		if partition != nil {
			t.Error("should return nil partition")
		}
		if err == nil {
			t.Error("should not return nil error")
		}
		expected := fmt.Sprintf("Invalid partition")
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("Valid partition", func(t *testing.T) {
		b, err := ioutil.ReadFile(mbrPartitionFile)
		if err != nil {
			t.Fatalf("Unable to read test fixture file %s: %v", mbrPartitionFile, err)
		}
		partition, err := partitionFromBytes(b)
		if partition == nil {
			t.Error("should not return nil partition")
		}
		if err != nil {
			t.Errorf("returned non-nil error: %v", err)
		}
		// check out data
		expected := Partition{
			Bootable:      false,
			StartHead:     0x20,
			StartSector:   0x21,
			StartCylinder: 0x00,
			Type:          Linux,
			EndHead:       0x31,
			EndSector:     0x18,
			EndCylinder:   0x00,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		if partition == nil && &expected != nil || !partition.Equal(&expected) {
			t.Log(b)
			t.Errorf("actual partition was %v instead of expected %v", *partition, expected)
		}

	})
}

func TestToBytes(t *testing.T) {
	t.Run("Valid partition", func(t *testing.T) {
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		b, err := partition.toBytes()
		if b == nil {
			t.Error("should not return nil bytes")
		}
		if err != nil {
			t.Error("should return nil error")
		}
		expected, err := ioutil.ReadFile(mbrPartitionFile)
		if err != nil {
			t.Fatalf("Unable to read test fixture file %s: %v", mbrPartitionFile, err)
		}
		if !PartitionEqualBytes(expected, b) {
			t.Errorf("returned byte %v instead of expected %v", b, expected)
		}
	})
}

func TestReadContents(t *testing.T) {
	t.Run("error reading file", func(t *testing.T) {
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		var b bytes.Buffer
		writer := bufio.NewWriter(&b)
		expected := "Error reading from file"
		f := &testhelper.FileImpl{
			Reader: func(b []byte, offset int64) (int, error) {
				return 0, fmt.Errorf(expected)
			},
		}
		read, err := partition.readContents(f, 512, 512, writer)
		if read != 0 {
			t.Errorf("Returned %d bytes read instead of 0", read)
		}
		if err == nil {
			t.Errorf("Returned nil error instead of actual errors")
		}
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("successful read", func(t *testing.T) {
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		var b bytes.Buffer
		writer := bufio.NewWriter(&b)
		size := 100
		b2 := make([]byte, size, size)
		rand.Read(b2)
		f := &testhelper.FileImpl{
			Reader: func(b []byte, offset int64) (int, error) {
				copy(b, b2)
				return size, io.EOF
			},
		}
		read, err := partition.readContents(f, 512, 512, writer)
		if read != int64(size) {
			t.Errorf("Returned %d bytes read instead of %d", read, size)
		}
		if err != nil {
			t.Errorf("Returned error instead of expected nil")
		}
		writer.Flush()
		if bytes.Compare(b.Bytes(), b2) != 0 {
			t.Errorf("Mismatched bytes data")
			t.Log(b.Bytes())
			t.Log(b2)
		}
	})
}

func TestWriteContents(t *testing.T) {
	t.Run("mismatched size", func(t *testing.T) {
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		var b bytes.Buffer
		reader := bufio.NewReader(&b)
		expected := "Write 0 bytes to partition "
		f := &testhelper.FileImpl{}
		written, err := partition.writeContents(f, 512, 512, reader)
		if written != 0 {
			t.Errorf("Returned %d bytes written instead of 0", written)
		}
		if err == nil {
			t.Errorf("Returned nil error instead of actual errors")
		}
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("error writing file", func(t *testing.T) {
		size := 512000
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		b := make([]byte, size, size)
		rand.Read(b)
		reader := bytes.NewReader(b)
		expected := "Error writing to file"
		f := &testhelper.FileImpl{
			Writer: func(b []byte, offset int64) (int, error) {
				return 0, fmt.Errorf(expected)
			},
		}
		written, err := partition.writeContents(f, 512, 512, reader)
		if written != 0 {
			t.Errorf("Returned %d bytes written instead of 0", written)
		}
		if err == nil {
			t.Errorf("Returned nil error instead of actual errors")
			return
		}
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("too large for partition", func(t *testing.T) {
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart / 512,
			Size:          1,
		}
		// make a byte array that is too big
		b := make([]byte, 2*512, 2*512)
		rand.Read(b)
		reader := bytes.NewReader(b)
		expected := "Requested to write at least"
		f := &testhelper.FileImpl{
			Writer: func(b []byte, offset int64) (int, error) {
				return len(b), nil
			},
		}
		written, err := partition.writeContents(f, 512, 512, reader)
		if written != 512 {
			t.Errorf("Returned %d bytes written instead of 512", written)
		}
		if err == nil {
			t.Errorf("Returned nil error instead of actual errors")
			return
		}
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})

	t.Run("successful write", func(t *testing.T) {
		size := 512000
		sectorSize := size / 512
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          uint32(sectorSize),
		}
		b := make([]byte, size, size)
		rand.Read(b)
		b2 := make([]byte, 0, size)
		reader := bytes.NewReader(b)
		f := &testhelper.FileImpl{
			Writer: func(b []byte, offset int64) (int, error) {
				b2 = append(b2, b...)
				return len(b), nil
			},
		}
		written, err := partition.writeContents(f, 512, 512, reader)
		if written != uint64(size) {
			t.Errorf("Returned %d bytes written instead of %d", written, size)
		}
		if err != nil {
			t.Errorf("Returned error instead of nil: %v", err)
			return
		}
		if bytes.Compare(b2, b) != 0 {
			t.Errorf("Bytes mismatch")
			t.Log(b)
			t.Log(b2)
		}
	})
}
<commit_msg>Updating per discussion.<commit_after>package mbr

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"fmt"
	"io"
	"io/ioutil"
	"strings"
	"testing"

	"github.com/diskfs/go-diskfs/testhelper"
)

const (
	mbrPartitionFile = "./testdata/mbr_partition.dat"
	partitionStart   = 2048
	partitionSize    = 20480
)

func TestFromBytes(t *testing.T) {
	t.Run("Short byte slice", func(t *testing.T) {
		b := make([]byte, partitionEntrySize-1, partitionEntrySize-1)
		rand.Read(b)
		partition, err := partitionFromBytes(b)
		if partition != nil {
			t.Error("should return nil partition")
		}
		if err == nil {
			t.Error("should not return nil error")
		}
		expected := fmt.Sprintf("Data for partition was %d bytes instead of expected %d", len(b), partitionEntrySize)
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("Long byte slice", func(t *testing.T) {
		b := make([]byte, partitionEntrySize+1, partitionEntrySize+1)
		rand.Read(b)
		partition, err := partitionFromBytes(b)
		if partition != nil {
			t.Error("should return nil partition")
		}
		if err == nil {
			t.Error("should not return nil error")
		}
		expected := fmt.Sprintf("Data for partition was %d bytes instead of expected %d", len(b), partitionEntrySize)
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("Invalid partition bootable code", func(t *testing.T) {
		b := make([]byte, partitionEntrySize, partitionEntrySize)
		rand.Read(b)
		b[0] = 0x67
		partition, err := partitionFromBytes(b)
		if partition != nil {
			t.Error("should return nil partition")
		}
		if err == nil {
			t.Error("should not return nil error")
		}
		expected := fmt.Sprintf("Invalid partition")
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("Valid partition", func(t *testing.T) {
		b, err := ioutil.ReadFile(mbrPartitionFile)
		if err != nil {
			t.Fatalf("Unable to read test fixture file %s: %v", mbrPartitionFile, err)
		}
		partition, err := partitionFromBytes(b)
		if partition == nil {
			t.Error("should not return nil partition")
		}
		if err != nil {
			t.Errorf("returned non-nil error: %v", err)
		}
		// check out data
		expected := Partition{
			Bootable:      false,
			StartHead:     0x20,
			StartSector:   0x21,
			StartCylinder: 0x00,
			Type:          Linux,
			EndHead:       0x31,
			EndSector:     0x18,
			EndCylinder:   0x00,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		if partition == nil && &expected != nil || !partition.Equal(&expected) {
			t.Log(b)
			t.Errorf("actual partition was %v instead of expected %v", *partition, expected)
		}

	})
}

func TestToBytes(t *testing.T) {
	t.Run("Valid partition", func(t *testing.T) {
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		b, err := partition.toBytes()
		if b == nil {
			t.Error("should not return nil bytes")
		}
		if err != nil {
			t.Error("should return nil error")
		}
		expected, err := ioutil.ReadFile(mbrPartitionFile)
		if err != nil {
			t.Fatalf("Unable to read test fixture file %s: %v", mbrPartitionFile, err)
		}
		if !PartitionEqualBytes(expected, b) {
			t.Errorf("returned byte %v instead of expected %v", b, expected)
		}
	})
}

func TestReadContents(t *testing.T) {
	t.Run("error reading file", func(t *testing.T) {
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		var b bytes.Buffer
		writer := bufio.NewWriter(&b)
		expected := "Error reading from file"
		f := &testhelper.FileImpl{
			Reader: func(b []byte, offset int64) (int, error) {
				return 0, fmt.Errorf(expected)
			},
		}
		read, err := partition.readContents(f, 512, 512, writer)
		if read != 0 {
			t.Errorf("Returned %d bytes read instead of 0", read)
		}
		if err == nil {
			t.Errorf("Returned nil error instead of actual errors")
		}
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("successful read", func(t *testing.T) {
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		var b bytes.Buffer
		writer := bufio.NewWriter(&b)
		size := 100
		b2 := make([]byte, size, size)
		rand.Read(b2)
		f := &testhelper.FileImpl{
			Reader: func(b []byte, offset int64) (int, error) {
				copy(b, b2)
				return size, io.EOF
			},
		}
		read, err := partition.readContents(f, 512, 512, writer)
		if read != int64(size) {
			t.Errorf("Returned %d bytes read instead of %d", read, size)
		}
		if err != nil {
			t.Errorf("Returned error instead of expected nil")
		}
		writer.Flush()
		if bytes.Compare(b.Bytes(), b2) != 0 {
			t.Errorf("Mismatched bytes data")
			t.Log(b.Bytes())
			t.Log(b2)
		}
	})
}

func TestWriteContents(t *testing.T) {
	t.Run("mismatched size", func(t *testing.T) {
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		var b bytes.Buffer
		reader := bufio.NewReader(&b)
		expected := "Write 0 bytes to partition "
		f := &testhelper.FileImpl{}
		written, err := partition.writeContents(f, 512, 512, reader)
		if written != 0 {
			t.Errorf("Returned %d bytes written instead of 0", written)
		}
		if err == nil {
			t.Errorf("Returned nil error instead of actual errors")
		}
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("error writing file", func(t *testing.T) {
		size := 512000
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          partitionSize,
		}
		b := make([]byte, size, size)
		rand.Read(b)
		reader := bytes.NewReader(b)
		expected := "Error writing to file"
		f := &testhelper.FileImpl{
			Writer: func(b []byte, offset int64) (int, error) {
				return 0, fmt.Errorf(expected)
			},
		}
		written, err := partition.writeContents(f, 512, 512, reader)
		if written != 0 {
			t.Errorf("Returned %d bytes written instead of 0", written)
		}
		if err == nil {
			t.Errorf("Returned nil error instead of actual errors")
			return
		}
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})
	t.Run("too large for partition", func(t *testing.T) {
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          1,
		}
		// make a byte array that is too big
		b := make([]byte, 2*512, 2*512)
		rand.Read(b)
		reader := bytes.NewReader(b)
		expected := "Requested to write at least"
		f := &testhelper.FileImpl{
			Writer: func(b []byte, offset int64) (int, error) {
				return len(b), nil
			},
		}
		// We have a size of 1 sector, or 512 bytes, but are trying to write 2*512.
		// It should write the first and fail on the second so we expect an error,
		// along with 512 bytes successfully written
		written, err := partition.writeContents(f, 512, 512, reader)
		if written != 512 {
			t.Errorf("Returned %d bytes written instead of 512", written)
		}
		if err == nil {
			t.Errorf("Returned nil error instead of actual errors")
			return
		}
		if !strings.HasPrefix(err.Error(), expected) {
			t.Errorf("Error type %s instead of expected %s", err.Error(), expected)
		}
	})

	t.Run("successful write", func(t *testing.T) {
		size := 512000
		sectorSize := size / 512
		partition := Partition{
			Bootable:      false,
			StartHead:     0,
			StartSector:   2,
			StartCylinder: 0,
			Type:          Linux,
			EndHead:       0,
			EndSector:     2,
			EndCylinder:   0,
			Start:         partitionStart,
			Size:          uint32(sectorSize),
		}
		b := make([]byte, size, size)
		rand.Read(b)
		b2 := make([]byte, 0, size)
		reader := bytes.NewReader(b)
		f := &testhelper.FileImpl{
			Writer: func(b []byte, offset int64) (int, error) {
				b2 = append(b2, b...)
				return len(b), nil
			},
		}
		written, err := partition.writeContents(f, 512, 512, reader)
		if written != uint64(size) {
			t.Errorf("Returned %d bytes written instead of %d", written, size)
		}
		if err != nil {
			t.Errorf("Returned error instead of nil: %v", err)
			return
		}
		if bytes.Compare(b2, b) != 0 {
			t.Errorf("Bytes mismatch")
			t.Log(b)
			t.Log(b2)
		}
	})
}
<|endoftext|>


<commit_before>/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package e2e

import (
	"fmt"
	"strings"
	"testing"
	"time"

	testify "github.com/stretchr/testify/assert"
	"gotest.tools/v3/assert"

	. "github.com/docker/compose-cli/utils/e2e"
)

func TestRestart(t *testing.T) {
	c := NewParallelE2eCLI(t, binDir)
	const projectName = "e2e-restart"

	getServiceRegx := func(service string, status string) string {
		// match output with random spaces like:
		// e2e-start-stop_db_1      db                  running
		return fmt.Sprintf("%s_%s_1\\s+%s\\s+%s", projectName, service, service, status)
	}

	t.Run("Up a project", func(t *testing.T) {
		// This is just to ensure the containers do NOT exist
		c.RunDockerOrExitError("compose", "--project-name", projectName, "down")

		res := c.RunDockerOrExitError("compose", "-f", "./fixtures/restart-test/compose.yml", "--project-name", projectName, "up", "-d")
		assert.Assert(t, strings.Contains(res.Combined(), "Container e2e-restart_restart_1  Started"), res.Combined())

		// Give the time for it to exit
		time.Sleep(time.Second)

		res = c.RunDockerOrExitError("compose", "--project-name", projectName, "ps", "-a")
		testify.Regexp(t, getServiceRegx("restart", "exited"), res.Stdout())

		_ = c.RunDockerOrExitError("compose", "-f", "./fixtures/restart-test/compose.yml", "--project-name", projectName, "restart")

		// Give the same time but it must NOT exit
		time.Sleep(time.Second)

		res = c.RunDockerOrExitError("compose", "--project-name", projectName, "ps")
		testify.Regexp(t, getServiceRegx("restart", "running"), res.Stdout())

		// Clean up
		c.RunDockerOrExitError("compose", "--project-name", projectName, "down")
	})
}
<commit_msg>Removed race in e2e test<commit_after>/*
   Copyright 2020 Docker Compose CLI authors

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

package e2e

import (
	"fmt"
	"strings"
	"testing"
	"time"

	testify "github.com/stretchr/testify/assert"
	"gotest.tools/v3/assert"

	. "github.com/docker/compose-cli/utils/e2e"
)

func TestRestart(t *testing.T) {
	c := NewParallelE2eCLI(t, binDir)
	const projectName = "e2e-restart"

	getServiceRegx := func(service string, status string) string {
		// match output with random spaces like:
		// e2e-start-stop_db_1      db                  running
		return fmt.Sprintf("%s_%s_1\\s+%s\\s+%s", projectName, service, service, status)
	}

	t.Run("Up a project", func(t *testing.T) {
		// This is just to ensure the containers do NOT exist
		c.RunDockerOrExitError("compose", "--project-name", projectName, "down")

		res := c.RunDockerOrExitError("compose", "-f", "./fixtures/restart-test/compose.yml", "--project-name", projectName, "up", "-d")
		assert.Assert(t, strings.Contains(res.Combined(), "Container e2e-restart_restart_1  Started"), res.Combined())

		c.WaitForCmdResult(c.NewDockerCmd("compose", "--project-name", projectName, "ps", "-a", "--format", "json"),
			StdoutContains(`"State":"exited"`),
			10*time.Second, 1*time.Second)

		res = c.RunDockerOrExitError("compose", "--project-name", projectName, "ps", "-a")
		testify.Regexp(t, getServiceRegx("restart", "exited"), res.Stdout())

		_ = c.RunDockerOrExitError("compose", "-f", "./fixtures/restart-test/compose.yml", "--project-name", projectName, "restart")

		// Give the same time but it must NOT exit
		time.Sleep(time.Second)

		res = c.RunDockerOrExitError("compose", "--project-name", projectName, "ps")
		testify.Regexp(t, getServiceRegx("restart", "running"), res.Stdout())

		// Clean up
		c.RunDockerOrExitError("compose", "--project-name", projectName, "down")
	})
}
<|endoftext|>


<commit_before>// Migrate from bolt to badger
package migrate_cache

import (
	"fmt"
	"time"

	"github.com/boltdb/bolt"
	"github.com/minotar/imgd/pkg/cache"
	"github.com/minotar/imgd/pkg/cache/badger_cache"
	"github.com/minotar/imgd/pkg/cache/bolt_cache"
	store_expiry "github.com/minotar/imgd/pkg/cache/util/expiry/store"
)

const MIN_RECACHE_TTL = time.Duration(1) * time.Minute

type MigrateCacheConfig struct {
	cache.CacheConfig
	bolt_cache.BoltCacheConfig
	badger_cache.BadgerCacheConfig
	OldCache *bolt_cache.BoltCache
	NewCache *badger_cache.BadgerCache
}

type MigrateCache struct {
	*MigrateCacheConfig
}

var _ cache.Cache = new(MigrateCache)

func NewMigrateCache(cfg *MigrateCacheConfig) (*MigrateCache, error) {
	cfg.Logger = cfg.Logger.With(
		"cacheName", cfg.Name,
		"cacheType", "MigrateCache",
		"cacheParent", "MigrateCache", // The caches will overwrite the cacheType
	)
	cfg.Logger.Infof("initializing MigrateCache")
	mc := &MigrateCache{MigrateCacheConfig: cfg}

	cfg.BoltCacheConfig.Logger = cfg.Logger
	cfg.BadgerCacheConfig.Logger = cfg.Logger

	bolt, err := bolt_cache.NewBoltCache(&cfg.BoltCacheConfig)
	if err != nil {
		return nil, err
	}

	badger, err := badger_cache.NewBadgerCache(&cfg.BadgerCacheConfig)
	if err != nil {
		return nil, err
	}

	mc.OldCache = bolt
	mc.NewCache = badger
	cfg.Logger.Infof("initialized MigrateCache \"%s\"", mc.Name())
	return mc, nil
}

func (mc *MigrateCache) Name() string {
	return mc.CacheConfig.Name
}

func (mc *MigrateCache) Insert(key string, value []byte) error {
	return mc.NewCache.Insert(key, value)
}

func (mc *MigrateCache) InsertTTL(key string, value []byte, ttl time.Duration) error {
	return mc.NewCache.InsertTTL(key, value, ttl)
}

func (mc *MigrateCache) Retrieve(key string) ([]byte, error) {
	var errors []error
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Retrieving \"%s\" from cache %d \"%s\"", key, i, c.Name())

		value, err := c.Retrieve(key)
		if err == cache.ErrNotFound {
			// errors logic at end handles ErrNotFound
			continue
		} else if err != nil {
			// This is a cache related error (vs. a missing key)
			mc.Logger.Errorf("Error retrieving key \"%s\" from cache %d (%s): %s", key, i, c.Name(), err)
			errors = append(errors, err)
			continue
		}

		return value, nil
	}

	// If we had a genuine error, `errors` would be populated, otherwise, it must just be ErrNotFound
	if errors != nil {
		return nil, fmt.Errorf("error(s) retrieving \"%s\" from cache(s): %+v", key, errors)
	}
	return nil, cache.ErrNotFound
}

// Probably won't be used too much
func (mc *MigrateCache) TTL(key string) (time.Duration, error) {
	var errors []error
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Getting TTL of key \"%s\" from cache %d (%s)", key, i, c.Name())

		ttl, err := c.TTL(key)
		if err == cache.ErrNoExpiry {
			// Record has no expiry - we trust the first cache with this response
			// It's important the cache signified between NoExpiry vs. NotFound!
			return ttl, err
		} else if err != nil {
			// Todo: Probably just print here?
			errors = append(errors, err)
			continue
		}
		return ttl, err
	}

	return 0, errors[len(errors)-1]
}

func (mc *MigrateCache) Remove(key string) error {
	var errors []error
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Removing key \"%s\" from cache %d (%s)", key, i, c.Name())

		err := c.Remove(key)
		if err != nil {
			mc.Logger.Errorf("Error removing key \"%s\" from cache %d (%s): %s", key, i, c.Name(), err)
			errors = append(errors, err)
		}
	}

	if errors != nil {
		return fmt.Errorf("error(s) removing \"%s\" from cache(s): %+v", key, errors)
	}
	return nil
}

func (mc *MigrateCache) Flush() error {
	var errors []error
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Flushing cache %d (%s)", i, c.Name())

		err := c.Flush()
		if err != nil {
			mc.Logger.Errorf("Error flushing cache %d (%s): %s", i, c.Name(), err)
			errors = append(errors, err)
		}
	}
	if errors != nil {
		return fmt.Errorf("error(s) flushing cache(s): %+v", errors)
	}
	return nil
}

func (mc *MigrateCache) Len() uint {
	var maxLen uint
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Getting length of cache %d (%s)", i, c.Name())

		cacheLen := c.Len()
		mc.Logger.Debugf("Length of cache %d (%s) is %d", i, c.Name(), cacheLen)
		if cacheLen > maxLen {
			maxLen = cacheLen
		}
	}
	return maxLen
}

func (mc *MigrateCache) Size() uint64 {
	var maxSize uint64
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Getting size of cache %d (%s)", i, c.Name())

		cacheSize := c.Size()
		mc.Logger.Debugf("Size of cache %d (%s) is %d", i, c.Name(), cacheSize)
		if cacheSize > maxSize {
			maxSize = cacheSize
		}
	}
	return maxSize
}

func (mc *MigrateCache) Migrate() {
	var scannedCount, errorCount int

	dbLength := int(mc.OldCache.Len())
	logger := mc.Logger.With("boltLength", dbLength)
	logger.Info("Starting migration from Bolt -> Badger")
	start := time.Now()

	mc.OldCache.DB.View(func(tx *bolt.Tx) error {
		c := tx.Bucket([]byte(mc.OldCache.Bucket)).Cursor()

		for k, v := c.First(); k != nil; k, v = c.Next() {
			scannedCount++
			key := string(k)

			present, _ := mc.NewCache.Key(key)
			if !present {
				bse := store_expiry.DecodeStoreEntry(k, v)
				ttl, _ := bse.TTL(mc.OldCache.StoreExpiry.Clock.Now())
				err := mc.NewCache.InsertTTL(key, bse.Value, ttl)
				if err != nil {
					errorCount++
				}
			}
		}

		return nil
	})

	dur := time.Since(start)
	logger = logger.With(
		"scannedCount", scannedCount,
		"errorCount", errorCount,
		"duration", dur,
	)
	logger.Info("Key Migration finished")
}

func (mc *MigrateCache) Start() {
	mc.Logger.Info("starting MigrateCache")
	mc.NewCache.Start()
	mc.Migrate()
}

func (mc *MigrateCache) Stop() {
	mc.NewCache.Stop()
}

func (mc *MigrateCache) Close() {
	mc.Logger.Debug("closing MigrateCache")
	mc.Stop()
	mc.OldCache.Close()
	mc.NewCache.Close()
}
<commit_msg>MigrateCache: background/go func for migration<commit_after>// Migrate from bolt to badger
package migrate_cache

import (
	"fmt"
	"time"

	"github.com/boltdb/bolt"
	"github.com/minotar/imgd/pkg/cache"
	"github.com/minotar/imgd/pkg/cache/badger_cache"
	"github.com/minotar/imgd/pkg/cache/bolt_cache"
	store_expiry "github.com/minotar/imgd/pkg/cache/util/expiry/store"
)

const MIN_RECACHE_TTL = time.Duration(1) * time.Minute

type MigrateCacheConfig struct {
	cache.CacheConfig
	bolt_cache.BoltCacheConfig
	badger_cache.BadgerCacheConfig
	OldCache *bolt_cache.BoltCache
	NewCache *badger_cache.BadgerCache
}

type MigrateCache struct {
	*MigrateCacheConfig
}

var _ cache.Cache = new(MigrateCache)

func NewMigrateCache(cfg *MigrateCacheConfig) (*MigrateCache, error) {
	cfg.Logger = cfg.Logger.With(
		"cacheName", cfg.Name,
		"cacheType", "MigrateCache",
		"cacheParent", "MigrateCache", // The caches will overwrite the cacheType
	)
	cfg.Logger.Infof("initializing MigrateCache")
	mc := &MigrateCache{MigrateCacheConfig: cfg}

	cfg.BoltCacheConfig.Logger = cfg.Logger
	cfg.BadgerCacheConfig.Logger = cfg.Logger

	bolt, err := bolt_cache.NewBoltCache(&cfg.BoltCacheConfig)
	if err != nil {
		return nil, err
	}

	badger, err := badger_cache.NewBadgerCache(&cfg.BadgerCacheConfig)
	if err != nil {
		return nil, err
	}

	mc.OldCache = bolt
	mc.NewCache = badger
	cfg.Logger.Infof("initialized MigrateCache \"%s\"", mc.Name())
	return mc, nil
}

func (mc *MigrateCache) Name() string {
	return mc.CacheConfig.Name
}

func (mc *MigrateCache) Insert(key string, value []byte) error {
	return mc.NewCache.Insert(key, value)
}

func (mc *MigrateCache) InsertTTL(key string, value []byte, ttl time.Duration) error {
	return mc.NewCache.InsertTTL(key, value, ttl)
}

func (mc *MigrateCache) Retrieve(key string) ([]byte, error) {
	var errors []error
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Retrieving \"%s\" from cache %d \"%s\"", key, i, c.Name())

		value, err := c.Retrieve(key)
		if err == cache.ErrNotFound {
			// errors logic at end handles ErrNotFound
			continue
		} else if err != nil {
			// This is a cache related error (vs. a missing key)
			mc.Logger.Errorf("Error retrieving key \"%s\" from cache %d (%s): %s", key, i, c.Name(), err)
			errors = append(errors, err)
			continue
		}

		return value, nil
	}

	// If we had a genuine error, `errors` would be populated, otherwise, it must just be ErrNotFound
	if errors != nil {
		return nil, fmt.Errorf("error(s) retrieving \"%s\" from cache(s): %+v", key, errors)
	}
	return nil, cache.ErrNotFound
}

// Probably won't be used too much
func (mc *MigrateCache) TTL(key string) (time.Duration, error) {
	var errors []error
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Getting TTL of key \"%s\" from cache %d (%s)", key, i, c.Name())

		ttl, err := c.TTL(key)
		if err == cache.ErrNoExpiry {
			// Record has no expiry - we trust the first cache with this response
			// It's important the cache signified between NoExpiry vs. NotFound!
			return ttl, err
		} else if err != nil {
			// Todo: Probably just print here?
			errors = append(errors, err)
			continue
		}
		return ttl, err
	}

	return 0, errors[len(errors)-1]
}

func (mc *MigrateCache) Remove(key string) error {
	var errors []error
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Removing key \"%s\" from cache %d (%s)", key, i, c.Name())

		err := c.Remove(key)
		if err != nil {
			mc.Logger.Errorf("Error removing key \"%s\" from cache %d (%s): %s", key, i, c.Name(), err)
			errors = append(errors, err)
		}
	}

	if errors != nil {
		return fmt.Errorf("error(s) removing \"%s\" from cache(s): %+v", key, errors)
	}
	return nil
}

func (mc *MigrateCache) Flush() error {
	var errors []error
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Flushing cache %d (%s)", i, c.Name())

		err := c.Flush()
		if err != nil {
			mc.Logger.Errorf("Error flushing cache %d (%s): %s", i, c.Name(), err)
			errors = append(errors, err)
		}
	}
	if errors != nil {
		return fmt.Errorf("error(s) flushing cache(s): %+v", errors)
	}
	return nil
}

func (mc *MigrateCache) Len() uint {
	var maxLen uint
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Getting length of cache %d (%s)", i, c.Name())

		cacheLen := c.Len()
		mc.Logger.Debugf("Length of cache %d (%s) is %d", i, c.Name(), cacheLen)
		if cacheLen > maxLen {
			maxLen = cacheLen
		}
	}
	return maxLen
}

func (mc *MigrateCache) Size() uint64 {
	var maxSize uint64
	for i, c := range []cache.Cache{mc.NewCache, mc.OldCache} {
		mc.Logger.Debugf("Getting size of cache %d (%s)", i, c.Name())

		cacheSize := c.Size()
		mc.Logger.Debugf("Size of cache %d (%s) is %d", i, c.Name(), cacheSize)
		if cacheSize > maxSize {
			maxSize = cacheSize
		}
	}
	return maxSize
}

func (mc *MigrateCache) Migrate() {
	var scannedCount, errorCount int

	dbLength := int(mc.OldCache.Len())
	logger := mc.Logger.With("boltLength", dbLength)
	logger.Info("Starting migration from Bolt -> Badger")
	start := time.Now()

	mc.OldCache.DB.View(func(tx *bolt.Tx) error {
		c := tx.Bucket([]byte(mc.OldCache.Bucket)).Cursor()

		for k, v := c.First(); k != nil; k, v = c.Next() {
			scannedCount++
			key := string(k)

			present, _ := mc.NewCache.Key(key)
			if !present {
				bse := store_expiry.DecodeStoreEntry(k, v)
				ttl, _ := bse.TTL(mc.OldCache.StoreExpiry.Clock.Now())
				err := mc.NewCache.InsertTTL(key, bse.Value, ttl)
				if err != nil {
					errorCount++
				}
			}
		}

		return nil
	})

	dur := time.Since(start)
	logger = logger.With(
		"scannedCount", scannedCount,
		"errorCount", errorCount,
		"duration", dur,
	)
	logger.Info("Key Migration finished")
}

func (mc *MigrateCache) Start() {
	mc.Logger.Info("starting MigrateCache")
	mc.NewCache.Start()
	go mc.Migrate()
}

func (mc *MigrateCache) Stop() {
	mc.NewCache.Stop()
}

func (mc *MigrateCache) Close() {
	mc.Logger.Debug("closing MigrateCache")
	mc.Stop()
	mc.OldCache.Close()
	mc.NewCache.Close()
}
<|endoftext|>