<commit_before>package gore

import (
	"errors"
	"math"
	"regexp"
	"strings"
	"sync"
	"time"
)

// Sentinel is a special Redis process that monitors other Redis instances,
// does fail-over, notifies client status of all monitored instances.
type Sentinel struct {
	servers   []string
	conn      *Conn
	subConn   *Conn // A dedicated connection for pubsub
	subs      *Subscriptions
	mutex     *sync.Mutex
	state     int
	instances map[string]*instance
}

// NewSentinel returns new Sentinel
func NewSentinel() *Sentinel {
	return &Sentinel{
		mutex:     &sync.Mutex{},
		state:     connStateNotConnected,
		instances: make(map[string]*instance),
	}
}

// AddServer adds new sentinel servers. Only one sentinel server is active
// at any time. If this server fails, gore will connect to other sentinel
// servers immediately.
//
// AddServer can be called at anytime, to add new server on the fly.
// In production environment, you should always have at least 3 sentinel
// servers up and running.
func (s *Sentinel) AddServer(addresses ...string) {
	s.servers = append(s.servers, addresses...)
}

// Dial connects to one sentinel server in the list. If it fails to connect,
// it moves to the next on the list. If all servers cannot be connected,
// Init return error.
func (s *Sentinel) Dial() (err error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if s.state != connStateNotConnected {
		return nil
	}
	return s.connect()
}

// Close gracefully closes the sentinel and all monitored connections
func (s *Sentinel) Close() {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.close()
}

// Ping pings all sentinel server, and returns error if it does not
// receive pong. This method can be called before Dial()
func (s *Sentinel) Ping() error {
	if len(s.servers) == 0 {
		return ErrNotConnected
	}
	conns := []*Conn{}
	defer func() {
		for _, conn := range conns {
			conn.Close()
		}
	}()
	for _, addr := range s.servers {
		conn, err := Dial(addr)
		if err != nil {
			return err
		}
		conns = append(conns, conn)
	}
	cmd := NewCommand("PING")
	for _, conn := range conns {
		_, err := cmd.Run(conn)
		if err != nil {
			return err
		}
	}
	return nil
}

// AddInstance registers a Redis instance with all sentinels server.
// This method can be called before Dial(). The quorum number is
// determined by the number of sentinel servers. If this method returns
// error, this may leave the sentinels data not synchronized.
func (s *Sentinel) AddInstance(name, host, port string) (err error) {
	if len(s.servers) == 0 {
		return ErrNotConnected
	}
	conns := []*Conn{}
	defer func() {
		for _, conn := range conns {
			conn.Close()
		}
	}()
	for _, addr := range s.servers {
		conn, err := Dial(addr)
		if err != nil {
			return err
		}
		conns = append(conns, conn)
	}
	quorum := math.Ceil(float64(len(s.servers)) / 2)
	cmd := NewCommand("SENTINEL", "MONITOR", name, host, port, quorum)
	for _, conn := range conns {
		rep, err := cmd.Run(conn)
		if err != nil {
			return err
		}
		if !rep.IsOk() {
			errMessage, _ := rep.Error()
			return errors.New(errMessage)
		}
	}
	return nil
}

// RemoveInstance removes a Redis instance from all sentinels server.
// Like AddInstance, this method can be called before Dial()
func (s *Sentinel) RemoveInstance(name string) error {
	if len(s.servers) == 0 {
		return ErrNotConnected
	}
	conns := []*Conn{}
	defer func() {
		for _, conn := range conns {
			conn.Close()
		}
	}()
	for _, addr := range s.servers {
		conn, err := Dial(addr)
		if err != nil {
			return err
		}
		conns = append(conns, conn)
	}
	cmd := NewCommand("SENTINEL", "REMOVE", name)
	for _, conn := range conns {
		_, err := cmd.Run(conn)
		if err != nil {
			return err
		}
	}
	return nil
}

// SetOption sets option for a Redis instance
// This method can be called before Dial()
func (s *Sentinel) SetOption(name, option, value string) error {
	if len(s.servers) == 0 {
		return ErrNotConnected
	}
	conns := []*Conn{}
	defer func() {
		for _, conn := range conns {
			conn.Close()
		}
	}()
	for _, addr := range s.servers {
		conn, err := Dial(addr)
		if err != nil {
			return err
		}
		conns = append(conns, conn)
	}
	cmd := NewCommand("SENTINEL", "SET", name, option, value)
	for _, conn := range conns {
		rep, err := cmd.Run(conn)
		if err != nil {
			return err
		}
		if !rep.IsOk() {
			errMessage, _ := rep.Error()
			return errors.New(errMessage)
		}
	}
	return nil
}

// GetPool returns a pool of connection from a pool name.
// If the pool has not been retrieved before, gore will attempt to
// fetch the address from the sentinel server, and initialize connections
// with this address. The application should never call this function repeatedly
// to get the same pool, because internal locking can cause performance to drop.
// An error can be returned if the pool name is not monitored by the sentinel,
// or the redis server is currently dead, or the redis server cannot be connected
// (for example: firewall issues).
func (s *Sentinel) GetPool(name string) (*Pool, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if ins, ok := s.instances[name]; ok {
		return ins.pool, nil
	}
	rep, err := NewCommand("SENTINEL", "master", name).Run(s.conn)
	if err != nil {
		return nil, err
	}
	if rep.IsError() {
		return nil, ErrNil
	}
	master, err := rep.Map()
	if err != nil {
		return nil, err
	}
	flags := strings.Split(master["flags"], ",")
	for _, flag := range flags {
		if flag == "s_down" || flag == "o_down" {
			return nil, ErrNotConnected
		}
	}
	ins := &instance{
		name:    name,
		address: master["ip"] + ":" + master["port"],
		state:   connStateConnected,
		pool:    &Pool{sentinel: true},
	}
	err = ins.pool.Dial(ins.address)
	if err != nil {
		return nil, err
	}
	s.instances[name] = ins
	return ins.pool, nil
}

// GetCluster return a cluster monitored by the sentinel.
// The name of the cluster will determine name of Redis instances.
// For example, if the cluster name is "mycluster", the instances' name
// maybe "mycluster1", "mycluster2", ...
func (s *Sentinel) GetCluster(name string) (c *Cluster, err error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	rep, err := NewCommand("SENTINEL", "masters").Run(s.conn)
	if err != nil {
		return nil, err
	}
	replies, err := rep.Array()
	if err != nil {
		return nil, err
	}
	if len(replies) == 0 {
		return nil, ErrNoShard
	}
	instances := make(map[string]*instance)
	defer func() {
		if err != nil {
			for _, ins := range instances {
				ins.pool.Close()
			}
		}
	}()
	for _, r := range replies {
		master, err := r.Map()
		if err != nil {
			return nil, err
		}
		suffix := strings.TrimPrefix(master["name"], name)
		if !suffixRegex.MatchString(suffix) {
			continue
		}
		ins := &instance{
			name:    master["name"],
			address: master["ip"] + ":" + master["port"],
			state:   connStateConnected,
			pool:    &Pool{sentinel: true},
		}
		err = ins.pool.Dial(ins.address)
		if err != nil {
			return nil, err
		}
		instances[ins.name] = ins
	}
	c = NewCluster()
	c.sentinel = true
	for _, ins := range instances {
		s.instances[ins.name] = ins
		c.addresses = append(c.addresses, ins.address)
		c.shards = append(c.shards, ins.pool)
	}
	return c, nil
}

var suffixRegex = regexp.MustCompile("^\\d+$")

func (s *Sentinel) connect() (err error) {
	for i, server := range s.servers {
		s.conn, err = DialTimeout(server, time.Duration(Config.ConnectTimeout)*time.Second)
		if err != nil {
			continue
		}
		s.subConn, err = DialTimeout(server, time.Duration(Config.ConnectTimeout)*time.Second)
		if err != nil {
			s.conn.Close()
			continue
		}
		s.state = connStateConnected
		s.subs = NewSubscriptions(s.subConn)
		s.subs.throwError = true
		err = s.subs.Subscribe("+sdown", "-sdown", "+odown", "-odown", "+switch-master")
		if err != nil {
			s.close()
			continue
		}
		s.servers = append(s.servers[0:i], s.servers[i+1:]...)
		s.servers = append(s.servers, server)
		go s.monitor()
		return nil
	}
	return ErrNotConnected
}

func (s *Sentinel) close() {
	s.state = connStateNotConnected
	s.subs.Close()
	s.subConn.Close()
	s.conn.Close()
	for _, ins := range s.instances {
		ins.pool.Close()
	}
	s.instances = make(map[string]*instance)
}

func (s *Sentinel) fail() {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if s.state == connStateConnected {
		s.state = connStateNotConnected
		s.reconnect()
	}
}

func (s *Sentinel) reconnect() {
	s.subs.Close()
	s.subConn.Close()
	s.conn.Close()
	sleepTime := Config.ReconnectTime
	for {
		err := s.connect()
		if err == nil {
			break
		}
		time.Sleep(time.Duration(sleepTime) * time.Second)
		if sleepTime < 30 {
			sleepTime += 2
		}
	}
}

func (s *Sentinel) monitor() {
	for message := range s.subs.Message() {
		if message == nil {
			s.fail()
			return
		}
		s.mutex.Lock()
		ins := s.getInstanceFromMessage(message)
		if ins == nil {
			s.mutex.Unlock()
			continue
		}
		if message.Channel == "+sdown" || message.Channel == "+odown" {
			ins.down(message)
		} else if message.Channel == "-sdown" || message.Channel == "-odown" {
			ins.up(message)
		} else if message.Channel == "+switch-master" {
			ins.switchMaster(s)
		}
		s.mutex.Unlock()
	}
}

func (s *Sentinel) getInstanceFromMessage(message *Message) *instance {
	if message.Channel == "+sdown" || message.Channel == "+odown" ||
		message.Channel == "-sdown" || message.Channel == "-odown" {
		pieces := strings.Split(string(message.Message), " ")
		if len(pieces) < 2 || pieces[0] != "master" {
			return nil
		}
		return s.instances[pieces[1]]
	} else if message.Channel == "+switch-master" {
		pieces := strings.Split(string(message.Message), " ")
		if len(pieces) < 1 {
			return nil
		}
		return s.instances[pieces[0]]
	}
	return nil
}

func (s *Sentinel) getInstanceAddress(name string) string {
	for {
		rep, err := NewCommand("SENTINEL", "get-master-addr-by-name", name).Run(s.conn)
		if err != nil {
			time.Sleep(2 * time.Second)
			continue
		}
		if !rep.IsArray() {
			return ""
		}
		result := []string{}
		rep.Slice(&result)
		if len(result) != 2 {
			return ""
		}
		return result[0] + ":" + result[1]
	}
}

type instance struct {
	name    string
	address string
	sdown   bool
	odown   bool
	pool    *Pool
	state   int
}

func (ins *instance) down(message *Message) {
	if ins.state != connStateConnected {
		return
	}
	ins.state = connStateNotConnected
	if message.Channel == "+sdown" {
		ins.sdown = true
	} else if message.Channel == "+odown" {
		ins.odown = true
	}
	ins.pool.sentinelGonnaLetYouDown()
}

func (ins *instance) up(message *Message) {
	if ins.state == connStateConnected {
		return
	}
	if message.Channel == "-sdown" {
		ins.sdown = false
	} else if message.Channel == "-odown" {
		ins.odown = false
	}
	if !ins.sdown && !ins.odown {
		ins.pool.sentinelGonnaGiveYouUp()
		ins.state = connStateConnected
	}
}

func (ins *instance) switchMaster(s *Sentinel) {
	address := s.getInstanceAddress(ins.name)
	if address == "" {
		// WTF
		return
	}
	if ins.state == connStateConnected {
		ins.pool.sentinelGonnaLetYouDown()
	}
	ins.pool.address = address
	ins.pool.sentinelGonnaGiveYouUp()
	ins.state = connStateConnected
}
<commit_msg>Added lock when adding sentinel server<commit_after>package gore

import (
	"errors"
	"math"
	"regexp"
	"strings"
	"sync"
	"time"
)

// Sentinel is a special Redis process that monitors other Redis instances,
// does fail-over, notifies client status of all monitored instances.
type Sentinel struct {
	servers   []string
	conn      *Conn
	subConn   *Conn // A dedicated connection for pubsub
	subs      *Subscriptions
	mutex     *sync.Mutex
	state     int
	instances map[string]*instance
}

// NewSentinel returns new Sentinel
func NewSentinel() *Sentinel {
	return &Sentinel{
		mutex:     &sync.Mutex{},
		state:     connStateNotConnected,
		instances: make(map[string]*instance),
	}
}

// AddServer adds new sentinel servers. Only one sentinel server is active
// at any time. If this server fails, gore will connect to other sentinel
// servers immediately.
//
// AddServer can be called at anytime, to add new server on the fly.
// In production environment, you should always have at least 3 sentinel
// servers up and running.
func (s *Sentinel) AddServer(addresses ...string) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.servers = append(s.servers, addresses...)
}

// Dial connects to one sentinel server in the list. If it fails to connect,
// it moves to the next on the list. If all servers cannot be connected,
// Init return error.
func (s *Sentinel) Dial() (err error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if s.state != connStateNotConnected {
		return nil
	}
	return s.connect()
}

// Close gracefully closes the sentinel and all monitored connections
func (s *Sentinel) Close() {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	s.close()
}

// Ping pings all sentinel server, and returns error if it does not
// receive pong. This method can be called before Dial()
func (s *Sentinel) Ping() error {
	if len(s.servers) == 0 {
		return ErrNotConnected
	}
	conns := []*Conn{}
	defer func() {
		for _, conn := range conns {
			conn.Close()
		}
	}()
	for _, addr := range s.servers {
		conn, err := Dial(addr)
		if err != nil {
			return err
		}
		conns = append(conns, conn)
	}
	cmd := NewCommand("PING")
	for _, conn := range conns {
		_, err := cmd.Run(conn)
		if err != nil {
			return err
		}
	}
	return nil
}

// AddInstance registers a Redis instance with all sentinels server.
// This method can be called before Dial(). The quorum number is
// determined by the number of sentinel servers. If this method returns
// error, this may leave the sentinels data not synchronized.
func (s *Sentinel) AddInstance(name, host, port string) (err error) {
	if len(s.servers) == 0 {
		return ErrNotConnected
	}
	conns := []*Conn{}
	defer func() {
		for _, conn := range conns {
			conn.Close()
		}
	}()
	for _, addr := range s.servers {
		conn, err := Dial(addr)
		if err != nil {
			return err
		}
		conns = append(conns, conn)
	}
	quorum := math.Ceil(float64(len(s.servers)) / 2)
	cmd := NewCommand("SENTINEL", "MONITOR", name, host, port, quorum)
	for _, conn := range conns {
		rep, err := cmd.Run(conn)
		if err != nil {
			return err
		}
		if !rep.IsOk() {
			errMessage, _ := rep.Error()
			return errors.New(errMessage)
		}
	}
	return nil
}

// RemoveInstance removes a Redis instance from all sentinels server.
// Like AddInstance, this method can be called before Dial()
func (s *Sentinel) RemoveInstance(name string) error {
	if len(s.servers) == 0 {
		return ErrNotConnected
	}
	conns := []*Conn{}
	defer func() {
		for _, conn := range conns {
			conn.Close()
		}
	}()
	for _, addr := range s.servers {
		conn, err := Dial(addr)
		if err != nil {
			return err
		}
		conns = append(conns, conn)
	}
	cmd := NewCommand("SENTINEL", "REMOVE", name)
	for _, conn := range conns {
		_, err := cmd.Run(conn)
		if err != nil {
			return err
		}
	}
	return nil
}

// SetOption sets option for a Redis instance
// This method can be called before Dial()
func (s *Sentinel) SetOption(name, option, value string) error {
	if len(s.servers) == 0 {
		return ErrNotConnected
	}
	conns := []*Conn{}
	defer func() {
		for _, conn := range conns {
			conn.Close()
		}
	}()
	for _, addr := range s.servers {
		conn, err := Dial(addr)
		if err != nil {
			return err
		}
		conns = append(conns, conn)
	}
	cmd := NewCommand("SENTINEL", "SET", name, option, value)
	for _, conn := range conns {
		rep, err := cmd.Run(conn)
		if err != nil {
			return err
		}
		if !rep.IsOk() {
			errMessage, _ := rep.Error()
			return errors.New(errMessage)
		}
	}
	return nil
}

// GetPool returns a pool of connection from a pool name.
// If the pool has not been retrieved before, gore will attempt to
// fetch the address from the sentinel server, and initialize connections
// with this address. The application should never call this function repeatedly
// to get the same pool, because internal locking can cause performance to drop.
// An error can be returned if the pool name is not monitored by the sentinel,
// or the redis server is currently dead, or the redis server cannot be connected
// (for example: firewall issues).
func (s *Sentinel) GetPool(name string) (*Pool, error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if ins, ok := s.instances[name]; ok {
		return ins.pool, nil
	}
	rep, err := NewCommand("SENTINEL", "master", name).Run(s.conn)
	if err != nil {
		return nil, err
	}
	if rep.IsError() {
		return nil, ErrNil
	}
	master, err := rep.Map()
	if err != nil {
		return nil, err
	}
	flags := strings.Split(master["flags"], ",")
	for _, flag := range flags {
		if flag == "s_down" || flag == "o_down" {
			return nil, ErrNotConnected
		}
	}
	ins := &instance{
		name:    name,
		address: master["ip"] + ":" + master["port"],
		state:   connStateConnected,
		pool:    &Pool{sentinel: true},
	}
	err = ins.pool.Dial(ins.address)
	if err != nil {
		return nil, err
	}
	s.instances[name] = ins
	return ins.pool, nil
}

// GetCluster return a cluster monitored by the sentinel.
// The name of the cluster will determine name of Redis instances.
// For example, if the cluster name is "mycluster", the instances' name
// maybe "mycluster1", "mycluster2", ...
func (s *Sentinel) GetCluster(name string) (c *Cluster, err error) {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	rep, err := NewCommand("SENTINEL", "masters").Run(s.conn)
	if err != nil {
		return nil, err
	}
	replies, err := rep.Array()
	if err != nil {
		return nil, err
	}
	if len(replies) == 0 {
		return nil, ErrNoShard
	}
	instances := make(map[string]*instance)
	defer func() {
		if err != nil {
			for _, ins := range instances {
				ins.pool.Close()
			}
		}
	}()
	for _, r := range replies {
		master, err := r.Map()
		if err != nil {
			return nil, err
		}
		suffix := strings.TrimPrefix(master["name"], name)
		if !suffixRegex.MatchString(suffix) {
			continue
		}
		ins := &instance{
			name:    master["name"],
			address: master["ip"] + ":" + master["port"],
			state:   connStateConnected,
			pool:    &Pool{sentinel: true},
		}
		err = ins.pool.Dial(ins.address)
		if err != nil {
			return nil, err
		}
		instances[ins.name] = ins
	}
	c = NewCluster()
	c.sentinel = true
	for _, ins := range instances {
		s.instances[ins.name] = ins
		c.addresses = append(c.addresses, ins.address)
		c.shards = append(c.shards, ins.pool)
	}
	return c, nil
}

var suffixRegex = regexp.MustCompile("^\\d+$")

func (s *Sentinel) connect() (err error) {
	for i, server := range s.servers {
		s.conn, err = DialTimeout(server, time.Duration(Config.ConnectTimeout)*time.Second)
		if err != nil {
			continue
		}
		s.subConn, err = DialTimeout(server, time.Duration(Config.ConnectTimeout)*time.Second)
		if err != nil {
			s.conn.Close()
			continue
		}
		s.state = connStateConnected
		s.subs = NewSubscriptions(s.subConn)
		s.subs.throwError = true
		err = s.subs.Subscribe("+sdown", "-sdown", "+odown", "-odown", "+switch-master")
		if err != nil {
			s.close()
			continue
		}
		s.servers = append(s.servers[0:i], s.servers[i+1:]...)
		s.servers = append(s.servers, server)
		go s.monitor()
		return nil
	}
	return ErrNotConnected
}

func (s *Sentinel) close() {
	s.state = connStateNotConnected
	s.subs.Close()
	s.subConn.Close()
	s.conn.Close()
	for _, ins := range s.instances {
		ins.pool.Close()
	}
	s.instances = make(map[string]*instance)
}

func (s *Sentinel) fail() {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if s.state == connStateConnected {
		s.state = connStateNotConnected
		s.reconnect()
	}
}

func (s *Sentinel) reconnect() {
	s.subs.Close()
	s.subConn.Close()
	s.conn.Close()
	sleepTime := Config.ReconnectTime
	for {
		err := s.connect()
		if err == nil {
			break
		}
		time.Sleep(time.Duration(sleepTime) * time.Second)
		if sleepTime < 30 {
			sleepTime += 2
		}
	}
}

func (s *Sentinel) monitor() {
	for message := range s.subs.Message() {
		if message == nil {
			s.fail()
			return
		}
		s.mutex.Lock()
		ins := s.getInstanceFromMessage(message)
		if ins == nil {
			s.mutex.Unlock()
			continue
		}
		if message.Channel == "+sdown" || message.Channel == "+odown" {
			ins.down(message)
		} else if message.Channel == "-sdown" || message.Channel == "-odown" {
			ins.up(message)
		} else if message.Channel == "+switch-master" {
			ins.switchMaster(s)
		}
		s.mutex.Unlock()
	}
}

func (s *Sentinel) getInstanceFromMessage(message *Message) *instance {
	if message.Channel == "+sdown" || message.Channel == "+odown" ||
		message.Channel == "-sdown" || message.Channel == "-odown" {
		pieces := strings.Split(string(message.Message), " ")
		if len(pieces) < 2 || pieces[0] != "master" {
			return nil
		}
		return s.instances[pieces[1]]
	} else if message.Channel == "+switch-master" {
		pieces := strings.Split(string(message.Message), " ")
		if len(pieces) < 1 {
			return nil
		}
		return s.instances[pieces[0]]
	}
	return nil
}

func (s *Sentinel) getInstanceAddress(name string) string {
	for {
		rep, err := NewCommand("SENTINEL", "get-master-addr-by-name", name).Run(s.conn)
		if err != nil {
			time.Sleep(2 * time.Second)
			continue
		}
		if !rep.IsArray() {
			return ""
		}
		result := []string{}
		rep.Slice(&result)
		if len(result) != 2 {
			return ""
		}
		return result[0] + ":" + result[1]
	}
}

type instance struct {
	name    string
	address string
	sdown   bool
	odown   bool
	pool    *Pool
	state   int
}

func (ins *instance) down(message *Message) {
	if ins.state != connStateConnected {
		return
	}
	ins.state = connStateNotConnected
	if message.Channel == "+sdown" {
		ins.sdown = true
	} else if message.Channel == "+odown" {
		ins.odown = true
	}
	ins.pool.sentinelGonnaLetYouDown()
}

func (ins *instance) up(message *Message) {
	if ins.state == connStateConnected {
		return
	}
	if message.Channel == "-sdown" {
		ins.sdown = false
	} else if message.Channel == "-odown" {
		ins.odown = false
	}
	if !ins.sdown && !ins.odown {
		ins.pool.sentinelGonnaGiveYouUp()
		ins.state = connStateConnected
	}
}

func (ins *instance) switchMaster(s *Sentinel) {
	address := s.getInstanceAddress(ins.name)
	if address == "" {
		// WTF
		return
	}
	if ins.state == connStateConnected {
		ins.pool.sentinelGonnaLetYouDown()
	}
	ins.pool.address = address
	ins.pool.sentinelGonnaGiveYouUp()
	ins.state = connStateConnected
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Intel Corporation
//
// 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 main

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"testing"
)

func fcntl(fd int, cmd int, arg int) (val int, err error) {
	r0, _, errno := syscall.Syscall(
		syscall.SYS_FCNTL, uintptr(fd), uintptr(cmd), uintptr(arg))
	val = int(r0)
	if errno != 0 {
		err = errno
	}
	return
}

// FdLeakDetector detects fd leaks by letting its user take snapshots of the
// list of opened fds at key points in the profiled program and compare two
// snapshots.
type FdLeakDetector struct {
}

type FdInfo struct {
	Fd          int
	Flags       int
	CloseOnExec bool
	Text        string
}

func (info *FdInfo) dump(w io.Writer) {
	flags := make([]string, 5)

	if info.CloseOnExec {
		flags = append(flags, "cloexec")
	}

	// file status
	if info.Flags&syscall.O_APPEND != 0 {
		flags = append(flags, "append")
	}
	if info.Flags&syscall.O_NONBLOCK != 0 {
		flags = append(flags, "nonblock")
	}

	// acc mode
	if info.Flags&syscall.O_RDONLY != 0 {
		flags = append(flags, "read-only")
	}
	if info.Flags&syscall.O_RDWR != 0 {
		flags = append(flags, "read-write")
	}
	if info.Flags&syscall.O_WRONLY != 0 {
		flags = append(flags, "write-only")
	}
	if info.Flags&syscall.O_DSYNC != 0 {
		flags = append(flags, "dsync")
	}
	if info.Flags&syscall.O_RSYNC != 0 {
		flags = append(flags, "rsync")
	}
	if info.Flags&syscall.O_SYNC != 0 {
		flags = append(flags, "sync")
	}

	fmt.Fprintf(w, "  %d: %s (%s)\n", info.Fd, info.Text,
		strings.Join(flags, ""))
}

func (info *FdInfo) equal(other *FdInfo) bool {
	return reflect.DeepEqual(info, other)
}

type FdSnapshot struct {
	Fds []FdInfo
}

// ByFd implements sort.Interface for []FdInfo based on the Fd field.
type ByFd []FdInfo

func (a ByFd) Len() int           { return len(a) }
func (a ByFd) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByFd) Less(i, j int) bool { return a[i].Fd < a[j].Fd }

func newFdSnapshot() *FdSnapshot {
	return &FdSnapshot{}
}

func (snap *FdSnapshot) FdInfo(i int) *FdInfo {
	return &snap.Fds[i]
}

func (snap *FdSnapshot) dump(w io.Writer) {
	for _, info := range snap.Fds {
		info.dump(w)
	}
}

// NewFdLeadDetector creates a new FdLeakDetector
func NewFdLeadDetector() *FdLeakDetector {
	return &FdLeakDetector{}
}

const selfFdPath = "/proc/self/fd"

// Snapshot captures the list of opened file descriptors of the current
// process
func (d *FdLeakDetector) Snapshot() (snap *FdSnapshot, err error) {
	root, err := os.Open(selfFdPath)
	if err != nil {
		return nil, err
	}
	defer root.Close()

	infos, err := root.Readdir(-1)
	if err != nil {
		return nil, err
	}

	snap = newFdSnapshot()
	for _, info := range infos {
		fd, err := strconv.Atoi(info.Name())
		if err != nil {
			return nil, err
		}

		// Don't capture the fd used to list /proc/self/fd
		if fd == int(root.Fd()) {
			continue
		}

		flags, err := fcntl(fd, syscall.F_GETFL, 0)
		if err != nil {
			return nil, err
		}

		fdFlags, err := fcntl(fd, syscall.F_GETFD, 0)
		if err != nil {
			return nil, err
		}

		// readlink on /proc/self/fd gives nice textual information
		// about the fd
		rl, err := os.Readlink(fmt.Sprintf("%s/%d", selfFdPath, fd))
		if err != nil {
			return nil, err
		}

		snap.Fds = append(snap.Fds, FdInfo{
			Fd:          fd,
			Flags:       flags,
			CloseOnExec: fdFlags&syscall.FD_CLOEXEC != 0,
			Text:        rl,
		})
	}

	sort.Sort(ByFd(snap.Fds))

	return
}

// Compare writes to w the differences between the a and b snaphots.
// true will be returned if the two snapshots are equal, false otherwise.
func (d *FdLeakDetector) Compare(w io.Writer, a, b *FdSnapshot) bool {
	equal := true

	i := 0
	j := 0

	for i < len(a.Fds) && j < len(b.Fds) {
		aInfo := a.FdInfo(i)
		bInfo := b.FdInfo(j)

		if aInfo.Fd == bInfo.Fd {
			// File descriptor found in both snapshots
			equal = equal && aInfo.equal(bInfo)
			i++
			j++

			if !equal {
				fmt.Fprintf(w, "- fd %d\n", aInfo.Fd)
				aInfo.dump(w)
				fmt.Fprintf(w, "+ fd %d\n", bInfo.Fd)
				bInfo.dump(w)
			}
		} else if aInfo.Fd < bInfo.Fd {
			// File descriptor present in a but not in b
			equal = false
			i++

			fmt.Fprintf(w, "- fd %d\n", aInfo.Fd)
			aInfo.dump(w)
		} else {
			// File descriptor present in b but not in a
			equal = false
			j++

			fmt.Fprintf(w, "+ fd %d\n", bInfo.Fd)
			bInfo.dump(w)
		}
	}

	for ; i < len(a.Fds); i++ {
		// File descriptor present in a but not in b
		equal = false
		aInfo := a.FdInfo(i)

		fmt.Fprintf(w, "- fd %d\n", aInfo.Fd)
		aInfo.dump(w)
	}

	for ; j < len(b.Fds); j++ {
		// File descriptor present in b but not in a
		equal = false
		bInfo := b.FdInfo(j)

		fmt.Fprintf(w, "+ fd %d\n", bInfo.Fd)
		bInfo.dump(w)
	}

	return equal
}

func TestFdDetectorNoLeak(t *testing.T) {
	detector := NewFdLeadDetector()

	old, err := detector.Snapshot()
	if err != nil {
		t.Error(err)
	}

	new, err := detector.Snapshot()
	if err != nil {
		t.Error(err)
	}

	buffer := bytes.NewBuffer(nil)
	equal := detector.Compare(buffer, old, new)
	if buffer.Len() != 0 {
		fmt.Print(buffer.String())
		t.Fatal()
	}
	if !equal {
		fmt.Print(buffer.String())
		t.Fatal()
	}
}

func TestFdDetectorLeak(t *testing.T) {
	detector := NewFdLeadDetector()

	old, err := detector.Snapshot()
	if err != nil {
		t.Error(err)
	}

	_, err = os.Open("/dev/null")
	if err != nil {
		t.Error(err)
	}

	new, err := detector.Snapshot()
	if err != nil {
		t.Error(err)
	}

	buffer := bytes.NewBuffer(nil)
	equal := detector.Compare(buffer, old, new)
	if equal {
		fmt.Print(buffer.String())
		t.Fatal()
	}
}

func TestFdDetectorCompare(t *testing.T) {
	tests := []struct {
		old, new *FdSnapshot
		equal    bool
	}{
		// Same fds
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			equal: true,
		},

		// Same fd number, different flags
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDONLY,
						Text:  "/foo",
					},
				},
			},
			equal: false,
		},

		// Same fd number, different close on exec status
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:          1,
						Flags:       syscall.O_RDWR,
						CloseOnExec: true,
						Text:        "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			equal: false,
		},

		// Same fd number, different readlink
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/bar",
					},
				},
			},
			equal: false,
		},

		// old has more fds
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
				},
			},
			equal: false,
		},

		// new has more fds
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			equal: false,
		},

		// Same number of fds in the two snapshots, different fds at
		// index 0 but same fds at index 1. We used to have a bug
		// saying the two snapshots were identical.
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDONLY,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			equal: false,
		},
	}

	detector := NewFdLeadDetector()

	for i, test := range tests {
		buffer := bytes.NewBuffer(nil)

		equal := detector.Compare(buffer, test.old, test.new)
		if equal != test.equal {
			test.old.dump(os.Stderr)
			test.new.dump(os.Stderr)
			fmt.Print(buffer.String())
			t.Fatal(fmt.Sprintf("Failed test #%d", i))
		}
	}
}
<commit_msg>fdleak: Handle nil snapshots in Compare<commit_after>// Copyright (c) 2017 Intel Corporation
//
// 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 main

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"reflect"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"testing"
)

func fcntl(fd int, cmd int, arg int) (val int, err error) {
	r0, _, errno := syscall.Syscall(
		syscall.SYS_FCNTL, uintptr(fd), uintptr(cmd), uintptr(arg))
	val = int(r0)
	if errno != 0 {
		err = errno
	}
	return
}

// FdLeakDetector detects fd leaks by letting its user take snapshots of the
// list of opened fds at key points in the profiled program and compare two
// snapshots.
type FdLeakDetector struct {
}

type FdInfo struct {
	Fd          int
	Flags       int
	CloseOnExec bool
	Text        string
}

func (info *FdInfo) dump(w io.Writer) {
	flags := make([]string, 5)

	if info.CloseOnExec {
		flags = append(flags, "cloexec")
	}

	// file status
	if info.Flags&syscall.O_APPEND != 0 {
		flags = append(flags, "append")
	}
	if info.Flags&syscall.O_NONBLOCK != 0 {
		flags = append(flags, "nonblock")
	}

	// acc mode
	if info.Flags&syscall.O_RDONLY != 0 {
		flags = append(flags, "read-only")
	}
	if info.Flags&syscall.O_RDWR != 0 {
		flags = append(flags, "read-write")
	}
	if info.Flags&syscall.O_WRONLY != 0 {
		flags = append(flags, "write-only")
	}
	if info.Flags&syscall.O_DSYNC != 0 {
		flags = append(flags, "dsync")
	}
	if info.Flags&syscall.O_RSYNC != 0 {
		flags = append(flags, "rsync")
	}
	if info.Flags&syscall.O_SYNC != 0 {
		flags = append(flags, "sync")
	}

	fmt.Fprintf(w, "  %d: %s (%s)\n", info.Fd, info.Text,
		strings.Join(flags, ""))
}

func (info *FdInfo) equal(other *FdInfo) bool {
	return reflect.DeepEqual(info, other)
}

type FdSnapshot struct {
	Fds []FdInfo
}

// ByFd implements sort.Interface for []FdInfo based on the Fd field.
type ByFd []FdInfo

func (a ByFd) Len() int           { return len(a) }
func (a ByFd) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a ByFd) Less(i, j int) bool { return a[i].Fd < a[j].Fd }

func newFdSnapshot() *FdSnapshot {
	return &FdSnapshot{}
}

func (snap *FdSnapshot) FdInfo(i int) *FdInfo {
	return &snap.Fds[i]
}

func (snap *FdSnapshot) dump(w io.Writer) {
	for _, info := range snap.Fds {
		info.dump(w)
	}
}

// NewFdLeadDetector creates a new FdLeakDetector
func NewFdLeadDetector() *FdLeakDetector {
	return &FdLeakDetector{}
}

const selfFdPath = "/proc/self/fd"

// Snapshot captures the list of opened file descriptors of the current
// process
func (d *FdLeakDetector) Snapshot() (snap *FdSnapshot, err error) {
	root, err := os.Open(selfFdPath)
	if err != nil {
		return nil, err
	}
	defer root.Close()

	infos, err := root.Readdir(-1)
	if err != nil {
		return nil, err
	}

	snap = newFdSnapshot()
	for _, info := range infos {
		fd, err := strconv.Atoi(info.Name())
		if err != nil {
			return nil, err
		}

		// Don't capture the fd used to list /proc/self/fd
		if fd == int(root.Fd()) {
			continue
		}

		flags, err := fcntl(fd, syscall.F_GETFL, 0)
		if err != nil {
			return nil, err
		}

		fdFlags, err := fcntl(fd, syscall.F_GETFD, 0)
		if err != nil {
			return nil, err
		}

		// readlink on /proc/self/fd gives nice textual information
		// about the fd
		rl, err := os.Readlink(fmt.Sprintf("%s/%d", selfFdPath, fd))
		if err != nil {
			return nil, err
		}

		snap.Fds = append(snap.Fds, FdInfo{
			Fd:          fd,
			Flags:       flags,
			CloseOnExec: fdFlags&syscall.FD_CLOEXEC != 0,
			Text:        rl,
		})
	}

	sort.Sort(ByFd(snap.Fds))

	return
}

// Compare writes to w the differences between the a and b snaphots.
// true will be returned if the two snapshots are equal, false otherwise.
func (d *FdLeakDetector) Compare(w io.Writer, a, b *FdSnapshot) bool {
	if a == nil || b == nil {
		return a == b
	}

	equal := true

	i := 0
	j := 0

	for i < len(a.Fds) && j < len(b.Fds) {
		aInfo := a.FdInfo(i)
		bInfo := b.FdInfo(j)

		if aInfo.Fd == bInfo.Fd {
			// File descriptor found in both snapshots
			equal = equal && aInfo.equal(bInfo)
			i++
			j++

			if !equal {
				fmt.Fprintf(w, "- fd %d\n", aInfo.Fd)
				aInfo.dump(w)
				fmt.Fprintf(w, "+ fd %d\n", bInfo.Fd)
				bInfo.dump(w)
			}
		} else if aInfo.Fd < bInfo.Fd {
			// File descriptor present in a but not in b
			equal = false
			i++

			fmt.Fprintf(w, "- fd %d\n", aInfo.Fd)
			aInfo.dump(w)
		} else {
			// File descriptor present in b but not in a
			equal = false
			j++

			fmt.Fprintf(w, "+ fd %d\n", bInfo.Fd)
			bInfo.dump(w)
		}
	}

	for ; i < len(a.Fds); i++ {
		// File descriptor present in a but not in b
		equal = false
		aInfo := a.FdInfo(i)

		fmt.Fprintf(w, "- fd %d\n", aInfo.Fd)
		aInfo.dump(w)
	}

	for ; j < len(b.Fds); j++ {
		// File descriptor present in b but not in a
		equal = false
		bInfo := b.FdInfo(j)

		fmt.Fprintf(w, "+ fd %d\n", bInfo.Fd)
		bInfo.dump(w)
	}

	return equal
}

func TestFdDetectorNoLeak(t *testing.T) {
	detector := NewFdLeadDetector()

	old, err := detector.Snapshot()
	if err != nil {
		t.Error(err)
	}

	new, err := detector.Snapshot()
	if err != nil {
		t.Error(err)
	}

	buffer := bytes.NewBuffer(nil)
	equal := detector.Compare(buffer, old, new)
	if buffer.Len() != 0 {
		fmt.Print(buffer.String())
		t.Fatal()
	}
	if !equal {
		fmt.Print(buffer.String())
		t.Fatal()
	}
}

func TestFdDetectorLeak(t *testing.T) {
	detector := NewFdLeadDetector()

	old, err := detector.Snapshot()
	if err != nil {
		t.Error(err)
	}

	_, err = os.Open("/dev/null")
	if err != nil {
		t.Error(err)
	}

	new, err := detector.Snapshot()
	if err != nil {
		t.Error(err)
	}

	buffer := bytes.NewBuffer(nil)
	equal := detector.Compare(buffer, old, new)
	if equal {
		fmt.Print(buffer.String())
		t.Fatal()
	}
}

func TestFdDetectorCompare(t *testing.T) {
	tests := []struct {
		old, new *FdSnapshot
		equal    bool
	}{
		// handle nil arguments
		{old: nil, new: nil, equal: true},
		{old: &FdSnapshot{}, new: nil, equal: false},
		{old: nil, new: &FdSnapshot{}, equal: false},
		// Same fds
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			equal: true,
		},

		// Same fd number, different flags
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDONLY,
						Text:  "/foo",
					},
				},
			},
			equal: false,
		},

		// Same fd number, different close on exec status
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:          1,
						Flags:       syscall.O_RDWR,
						CloseOnExec: true,
						Text:        "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			equal: false,
		},

		// Same fd number, different readlink
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/bar",
					},
				},
			},
			equal: false,
		},

		// old has more fds
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
				},
			},
			equal: false,
		},

		// new has more fds
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			equal: false,
		},

		// Same number of fds in the two snapshots, different fds at
		// index 0 but same fds at index 1. We used to have a bug
		// saying the two snapshots were identical.
		{
			old: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDONLY,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			new: &FdSnapshot{
				Fds: []FdInfo{
					{
						Fd:    0,
						Flags: syscall.O_RDWR,
					},
					{
						Fd:    1,
						Flags: syscall.O_RDWR,
						Text:  "/foo",
					},
				},
			},
			equal: false,
		},
	}

	detector := NewFdLeadDetector()

	for i, test := range tests {
		buffer := bytes.NewBuffer(nil)

		equal := detector.Compare(buffer, test.old, test.new)
		if equal != test.equal {
			test.old.dump(os.Stderr)
			test.new.dump(os.Stderr)
			fmt.Print(buffer.String())
			t.Fatal(fmt.Sprintf("Failed test #%d", i))
		}
	}
}
<|endoftext|>


<commit_before>package file

import (
	"fmt"
	"strings"

	"github.com/antham/doc-hunt/util"
)

// Config represents a config line
type Config struct {
	Doc     Doc
	Sources []Source
}

// ListConfig return a config list
func ListConfig() (*[]Config, error) {
	configs := []Config{}

	rows, err := db.Query("select d.id, d.category, d.identifier, d.created_at, s.id, s.category, s.identifier, s.created_at, s.doc_id from docs d inner join sources s on s.doc_id = d.id order by d.created_at")

	if err != nil {
		return nil, err
	}

	for rows.Next() {
		doc := Doc{}
		source := Source{}

		err := rows.Scan(&doc.ID, &doc.Category, &doc.Identifier, &doc.CreatedAt, &source.ID, &source.Category, &source.Identifier, &source.CreatedAt, &source.DocID)

		if err != nil {
			return nil, err
		}

		if len(configs) == 0 || configs[len(configs)-1].Doc.ID != source.DocID {
			configs = append(configs, Config{
				Doc: doc,
			})
		}

		configs[len(configs)-1].Sources = append(configs[len(configs)-1].Sources, source)
	}

	return &configs, nil
}

// CreateConfig populates everything needed to create a new config entry
func CreateConfig(docIdentifier string, docCat DocCategory, folderSources []string, fileSources []string) error {
	doc := NewDoc(docIdentifier, docCat)
	InsertDoc(doc)

	for _, identifier := range folderSources {
		source := NewSource(doc, identifier, SFOLDER)
		err := InsertSource(source)

		if err != nil {
			return err
		}

		items, err := NewItems(util.ExtractFolderFiles(identifier), source)

		if err != nil {
			return err
		}

		err = InsertItems(items)

		if err != nil {
			return err
		}
	}

	for _, identifier := range fileSources {
		source := NewSource(doc, identifier, SFILE)
		err := InsertSource(source)

		if err != nil {
			return err
		}

		items, err := NewItems(&[]string{identifier}, source)

		if err != nil {
			return err
		}

		InsertItems(items)
	}

	return nil
}

// RemoveConfigs delete one or several config group
func RemoveConfigs(configs *[]Config) error {
	sourceIds := []string{}
	docIds := []string{}

	for _, config := range *configs {
		for _, source := range config.Sources {
			sourceIds = append(sourceIds, fmt.Sprintf(`"%s"`, source.ID))
		}
		docIds = append(docIds, fmt.Sprintf(`"%s"`, config.Doc.ID))
	}

	if len(sourceIds) > 0 {
		_, err := db.Exec(fmt.Sprintf("delete from sources where id in (%s)", strings.Join(sourceIds, ",")))

		if err != nil {
			return err
		}

		_, err = db.Exec(fmt.Sprintf("delete from docs where id in (%s)", strings.Join(docIds, ",")))

		if err != nil {
			return err
		}

		_, err = db.Exec(fmt.Sprintf("delete from items where source_id in (%s)", strings.Join(sourceIds, ",")))

		if err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>ref(file/config) : split function<commit_after>package file

import (
	"fmt"
	"strings"

	"github.com/antham/doc-hunt/util"
)

// Config represents a config line
type Config struct {
	Doc     Doc
	Sources []Source
}

func createFolderSource(doc *Doc, folderSources *[]string) error {
	for _, identifier := range *folderSources {
		source := NewSource(doc, identifier, SFOLDER)
		err := InsertSource(source)

		if err != nil {
			return err
		}

		files, err := util.ExtractFolderFiles(identifier)

		if err != nil {
			return err
		}

		items, err := NewItems(files, source)

		if err != nil {
			return err
		}

		err = InsertItems(items)

		if err != nil {
			return err
		}
	}

	return nil
}

func createFileSources(doc *Doc, fileSources *[]string) error {
	for _, identifier := range *fileSources {
		source := NewSource(doc, identifier, SFILE)
		err := InsertSource(source)

		if err != nil {
			return err
		}

		items, err := NewItems(&[]string{identifier}, source)

		if err != nil {
			return err
		}

		err = InsertItems(items)

		if err != nil {
			return err
		}
	}

	return nil
}

// ListConfig return a config list
func ListConfig() (*[]Config, error) {
	configs := []Config{}

	rows, err := db.Query("select d.id, d.category, d.identifier, d.created_at, s.id, s.category, s.identifier, s.created_at, s.doc_id from docs d inner join sources s on s.doc_id = d.id order by d.created_at")

	if err != nil {
		return nil, err
	}

	for rows.Next() {
		doc := Doc{}
		source := Source{}

		err := rows.Scan(&doc.ID, &doc.Category, &doc.Identifier, &doc.CreatedAt, &source.ID, &source.Category, &source.Identifier, &source.CreatedAt, &source.DocID)

		if err != nil {
			return nil, err
		}

		if len(configs) == 0 || configs[len(configs)-1].Doc.ID != source.DocID {
			configs = append(configs, Config{
				Doc: doc,
			})
		}

		configs[len(configs)-1].Sources = append(configs[len(configs)-1].Sources, source)
	}

	return &configs, nil
}

// CreateConfig populates everything needed to create a new config entry
func CreateConfig(docIdentifier string, docCat DocCategory, folderSources []string, fileSources []string) error {
	doc := NewDoc(docIdentifier, docCat)
	err := InsertDoc(doc)

	if err != nil {
		return err
	}

	err = createFolderSource(doc, &folderSources)

	if err != nil {
		return err
	}

	err = createFileSources(doc, &fileSources)

	return err
}

// RemoveConfigs delete one or several config group
func RemoveConfigs(configs *[]Config) error {
	sourceIds := []string{}
	docIds := []string{}

	for _, config := range *configs {
		for _, source := range config.Sources {
			sourceIds = append(sourceIds, fmt.Sprintf(`"%s"`, source.ID))
		}
		docIds = append(docIds, fmt.Sprintf(`"%s"`, config.Doc.ID))
	}

	if len(sourceIds) > 0 {
		_, err := db.Exec(fmt.Sprintf("delete from sources where id in (%s)", strings.Join(sourceIds, ",")))

		if err != nil {
			return err
		}

		_, err = db.Exec(fmt.Sprintf("delete from docs where id in (%s)", strings.Join(docIds, ",")))

		if err != nil {
			return err
		}

		_, err = db.Exec(fmt.Sprintf("delete from items where source_id in (%s)", strings.Join(sourceIds, ",")))

		if err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package find

import (
	"fmt"
	"regexp"
	"strings"

	"github.com/lomik/graphite-clickhouse/helper/clickhouse"
)

func HasWildcard(target string) bool {
	return strings.IndexAny(target, "[]{}*") > -1
}

func MakeWhere(target string, withLevel bool) (where string) {
	level := strings.Count(target, ".") + 1

	AND := func(exp string) {
		if where == "" {
			where = exp
		} else {
			where = fmt.Sprintf("(%s) AND %s", where, exp)
		}
	}

	if withLevel {
		where = fmt.Sprintf("Level = %d", level)
	}

	if target == "*" {
		return
	}

	// simple metric
	if !HasWildcard(target) {
		AND(fmt.Sprintf("Path = '%s'", clickhouse.Escape(target)))
		return
	}

	// before any wildcard symbol
	simplePrefix := target[:strings.IndexAny(target, "[]{}*")]

	if len(simplePrefix) > 0 {
		AND(fmt.Sprintf("Path LIKE '%s%%'", clickhouse.Escape(simplePrefix)))
	}

	// prefix search like "metric.name.xx*"
	if len(simplePrefix) == len(target)-1 && target[len(target)-1] == '*' {
		return
	}

	pattern := GlobToRegexp(target)
	AND(fmt.Sprintf("match(Path, '%s')", clickhouse.Escape(pattern)))

	return
}

func GlobToRegexp(g string) string {
	s := g
	s = strings.Replace(s, "*", "([^.]*?)", -1)
	s = strings.Replace(s, "{", "(", -1)
	s = strings.Replace(s, "}", ")", -1)
	s = strings.Replace(s, ",", "|", -1)
	return s
}

func RemoveExtraPrefix(prefix, query string) (string, string, error) {
	qs := strings.Split(query, ".")
	ps := strings.Split(prefix, ".")

	var i int
	for i = 0; i < len(qs) && i < len(ps); i++ {
		m, err := regexp.MatchString(GlobToRegexp(qs[i]), ps[i])
		if err != nil {
			return "", "", err
		}
		if !m { // not matched
			return "", "", nil
		}
	}

	if i < len(ps) {
		return strings.Join(ps[:i], "."), "", nil
	}

	return prefix, strings.Join(qs[i:], "."), nil
}
<commit_msg>regexp begin and end<commit_after>package find

import (
	"fmt"
	"regexp"
	"strings"

	"github.com/lomik/graphite-clickhouse/helper/clickhouse"
)

func HasWildcard(target string) bool {
	return strings.IndexAny(target, "[]{}*") > -1
}

func MakeWhere(target string, withLevel bool) (where string) {
	level := strings.Count(target, ".") + 1

	AND := func(exp string) {
		if where == "" {
			where = exp
		} else {
			where = fmt.Sprintf("(%s) AND %s", where, exp)
		}
	}

	if withLevel {
		where = fmt.Sprintf("Level = %d", level)
	}

	if target == "*" {
		return
	}

	// simple metric
	if !HasWildcard(target) {
		AND(fmt.Sprintf("Path = '%s'", clickhouse.Escape(target)))
		return
	}

	// before any wildcard symbol
	simplePrefix := target[:strings.IndexAny(target, "[]{}*")]

	if len(simplePrefix) > 0 {
		AND(fmt.Sprintf("Path LIKE '%s%%'", clickhouse.Escape(simplePrefix)))
	}

	// prefix search like "metric.name.xx*"
	if len(simplePrefix) == len(target)-1 && target[len(target)-1] == '*' {
		return
	}

	pattern := GlobToRegexp(target)
	AND(fmt.Sprintf("match(Path, '^%s$')", clickhouse.Escape(pattern)))

	return
}

func GlobToRegexp(g string) string {
	s := g
	s = strings.Replace(s, "*", "([^.]*?)", -1)
	s = strings.Replace(s, "{", "(", -1)
	s = strings.Replace(s, "}", ")", -1)
	s = strings.Replace(s, ",", "|", -1)
	return s
}

func RemoveExtraPrefix(prefix, query string) (string, string, error) {
	qs := strings.Split(query, ".")
	ps := strings.Split(prefix, ".")

	var i int
	for i = 0; i < len(qs) && i < len(ps); i++ {
		m, err := regexp.MatchString(GlobToRegexp(qs[i]), ps[i])
		if err != nil {
			return "", "", err
		}
		if !m { // not matched
			return "", "", nil
		}
	}

	if i < len(ps) {
		return strings.Join(ps[:i], "."), "", nil
	}

	return prefix, strings.Join(qs[i:], "."), nil
}
<|endoftext|>


<commit_before>package jsons

import (
	"io"
	"os"
)

// FileWriter is a wrapper around a Writer which writes JSONS data in a file.
type FileWriter struct {
	filename string

	f io.WriteCloser
	w Writer
}

// NewFileWriter returns a pointer on a new FileWriter which will write in the
// given filename. The file is truncated if it already exists.
func NewFileWriter(filename string) *FileWriter {
	// TODO support append?
	return &FileWriter{filename: filename}
}

// Open opens the underlying file for writing. This must be called before any
// Add or Close call.
func (fw *FileWriter) Open() error {
	f, err := os.Create(fw.filename)
	if err != nil {
		return err
	}
	fw.f = f
	fw.w = NewWriter(f)
	return nil
}

// Close closes the underlying file. You must call this when you're done adding
// values.
func (fw *FileWriter) Close() error {
	return fw.f.Close()
}

// Add encodes the given value as a JSON object and write it in the underlying
// file.
func (fw *FileWriter) Add(v interface{}) error {
	return fw.w.Add(v)
}
<commit_msg>Add FileWriter.AddAll<commit_after>package jsons

import (
	"io"
	"os"
)

// FileWriter is a wrapper around a Writer which writes JSONS data in a file.
type FileWriter struct {
	filename string

	f io.WriteCloser
	w Writer
}

// NewFileWriter returns a pointer on a new FileWriter which will write in the
// given filename. The file is truncated if it already exists.
func NewFileWriter(filename string) *FileWriter {
	// TODO support append?
	return &FileWriter{filename: filename}
}

// Open opens the underlying file for writing. This must be called before any
// Add or Close call.
func (fw *FileWriter) Open() error {
	f, err := os.Create(fw.filename)
	if err != nil {
		return err
	}
	fw.f = f
	fw.w = NewWriter(f)
	return nil
}

// Close closes the underlying file. You must call this when you're done adding
// values.
func (fw *FileWriter) Close() error {
	return fw.f.Close()
}

// Add encodes the given value as a JSON object and write it in the underlying
// file.
func (fw *FileWriter) Add(v interface{}) error {
	return fw.w.Add(v)
}

// AddAll is equivalent to calling Add on each of its arguments
func (fw *FileWriter) AddAll(args ...interface{}) error {
	return fw.w.AddAll(args)
}
<|endoftext|>


<commit_before>package gokvlite

import (
	"container/list"
	"encoding/binary"
	"errors"
	"io"
	"os"
)

const freeBlockSize = 1024

type fileHeaderData struct {
	Freeblock_start int64
	Data_start      int64
}

type blockListHeaderData struct {
	Next int64
	Size int64
}

type blockListArrayEntryData struct {
	//This is a single array entry
	Free  uint8
	Start int64
	Size  int64
}

type blockListInfo struct {
	//This is a wrapper for a single array entry
	Entry    *blockListArrayEntryData
	Location int64
}

func (info *blockListInfo) writeInfo(writer io.WriterAt) error {
	return writeTo(writer, info.Location, info.Entry)
}

func (info *blockListInfo) ReadData(reader io.ReaderAt) (*[]byte, error) {
	//Reads the data from the location specific in the info/entry
	if info.Entry.Free > 0 {
		return nil, errors.New("filemanager: ReadData: Info is free, unable to read")
	}
	data := make([]byte, info.Entry.Size)
	_, err := reader.ReadAt(data, info.Entry.Start)
	return &data, err
}

func (info *blockListInfo) WriteData(writer io.WriterAt, data []byte) error {
	//Writes the data to the location specified by the info/entry
	if info.Entry.Free == 0 {
		return errors.New("Info is free, unable to write")
	}
	if int64(binary.Size(data)) != info.Entry.Size {
		return errors.New("Size of data is incorrect for size of info.")
	}
	_, err := writer.WriteAt(data, info.Entry.Start)
	return err
}

type blockListManager struct {
	headerStart int64
	header      *blockListHeaderData
}

type blockListInterface struct {
	/* This handles multiple block lists, gets a free block, etc
	 */

	fileheader     *fileHeaderData
	file           *os.File
	Blocklists     list.List
	Freeblocks     list.List
	Freeentries    list.List
	BlockListInfos map[int64]*blockListInfo
}

func (bli *blockListInterface) getFileEnd() (int64, error) {
	err := bli.file.Sync()
	if err != nil {
		return 0, err
	}

	fi, err := bli.file.Stat()
	if err != nil {
		return 0, err
	}

	end := fi.Size()
	return end, err
}

func (bli *blockListInterface) makeNewBlockList() error {
	end, err := bli.getFileEnd()
	if err != nil {
		return err
	}
	manager, _, err := bli.newBlockList(bli.file, end, freeBlockSize)
	if err != nil {
		return err
	}

	el := bli.Blocklists.Back()
	if el != nil {
		bl, ok := el.Value.(*blockListManager)
		if !ok {
			return errors.New("Invalid type for blocklistmanager in makeNewBlockList")
		}

		bl.header.Next = manager.headerStart
		err = writeTo(bli.file, bl.headerStart, bl.header)
		if err != nil {
			return err
		}
	}

	bli.Blocklists.PushBack(manager)
	return nil
}

func (bli *blockListInterface) getFreeEntry() (*blockListInfo, error) {
	//gets a free entry and removes it from the free entry list
	el := bli.Freeentries.Front()
	if el == nil {
		err := bli.makeNewBlockList()
		if err != nil {
			return nil, err
		}

		el = bli.Freeentries.Front()
		if el == nil {
			return nil, errors.New("Unable to get new element after making new blocklist")
		}
	}

	info, ok := el.Value.(*blockListInfo)
	if !ok {
		return info, errors.New("Invalid type in Freeentries list")
	}
	bli.Freeentries.Remove(el)

	return info, nil
}

func (bli *blockListInterface) Resize(info *blockListInfo, size int64) (*blockListInfo, *blockListInfo, error) {
	//resizes the info block, makes a new block, returns the old resized one and the new one
	startingsize := info.Entry.Size
	if startingsize == size {
		return nil, info, nil
	}

	if size > startingsize {
		//Just get a new entry and disregard the existing one
		newinfo, err := bli.GetFree(size)
		if err != nil {
			return nil, nil, err
		}
		err = bli.SetFree(info)
		if err != nil {
			return nil, nil, err
		}
		return nil, newinfo, nil
	}

	info.Entry.Size = size

	err := info.writeInfo(bli.file)
	if err != nil {
		return nil, nil, err
	}

	newinfo, err := bli.getFreeEntry()
	if err != nil {
		return nil, nil, err
	}

	newinfo.Entry.Size = startingsize - size
	newinfo.Entry.Free = 1
	newinfo.Entry.Start = info.Entry.Start + size
	err = newinfo.writeInfo(bli.file)
	if err != nil {
		return nil, nil, err
	}
	bli.Freeblocks.PushBack(newinfo)

	return info, newinfo, nil
}

func (bli *blockListInterface) SetFree(info *blockListInfo) error {
	//Sets a blocklist as free
	info.Entry.Free = 1
	err := info.writeInfo(bli.file)
	if err != nil {
		return err
	}
	bli.Freeblocks.PushBack(info)
	return nil
}

func (bli *blockListInterface) GetFree(size int64) (*blockListInfo, error) {
	//find a block large enough and resize it if it exists. Returns it marked as not free
	for e := bli.Freeblocks.Front(); e != nil; e = e.Next() {
		info, ok := e.Value.(*blockListInfo)
		if !ok {
			return nil, errors.New("Incorrect type in Freeblocks")
		}
		switch {
		case info.Entry.Size == size:
			info.Entry.Free = 0
			bli.Freeblocks.Remove(e)
			return info, nil

		case info.Entry.Size > size:
			_, _, err := bli.Resize(info, size)
			if err != nil {
				return nil, err
			}
			info.Entry.Free = 0
			bli.Freeblocks.Remove(e)
			return info, nil
		}
	}

	//append to end of file
	info, err := bli.getFreeEntry()
	if err != nil {
		return nil, err
	}

	end, err := bli.getFileEnd()
	if err != nil {
		return nil, err
	}

	info.Entry.Start = end
	info.Entry.Size = size
	info.Entry.Free = 0

	//Allocate the data storage
	data := make([]byte, size)
	bli.file.WriteAt(data, end)

	err = info.writeInfo(bli.file)
	return info, err
}

func (bli *blockListInterface) newBlockList(w io.WriterAt, start int64, size int64) (*blockListManager, int64, error) {
	//This creates a single unlinked blocklist
	//Size is in number of items, not bytes
	var written int64

	manager := new(blockListManager)

	header := new(blockListHeaderData)
	header.Size = size
	header.Next = 0
	err := writeTo(w, start, header)

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

	written += int64(binary.Size(header))

	manager.header = header
	manager.headerStart = start

	for i := int64(0); i < size; i++ {
		entry := new(blockListArrayEntryData)
		entry.Free = 1
		info := blockListInfo{entry, written + start}
		err = info.writeInfo(w)
		if err != nil {
			return manager, written, err
		}
		bli.Freeentries.PushBack(&info)
		bli.BlockListInfos[info.Location] = &info
		written += int64(binary.Size(entry))

	}
	return manager, written, err
}

func (bli *blockListInterface) readBlockList(reader io.ReaderAt, start int64) (*blockListManager, error) {
	var read int64
	header := new(blockListHeaderData)
	err := readFrom(reader, start, header)
	if err != nil {
		return nil, err
	}

	blm := new(blockListManager)
	blm.header = header
	blm.headerStart = start

	read += int64(binary.Size(header))

	for i := int64(0); i < blm.header.Size; i++ {
		data := new(blockListArrayEntryData)
		err := readFrom(reader, start+read, data)
		if err != nil {
			return blm, err
		}

		info := blockListInfo{data, start + read}
		if data.Size == 0 {
			bli.Freeentries.PushBack(&info)
		} else if data.Free > 0 {
			bli.Freeentries.PushBack(&info)
		}
		bli.BlockListInfos[info.Location] = &info

		read += int64(binary.Size(data))
	}
	return blm, nil
}

func readFrom(r io.ReaderAt, start int64, data interface{}) error {
	//This function reads any of the structs above and returns it
	sr := io.NewSectionReader(r, start, int64(binary.Size(data)))
	return binary.Read(sr, binary.LittleEndian, data)
}

func writeTo(w io.WriterAt, start int64, data interface{}) error {
	//This function writes any of the structs above and returns it
	sw := newSectionWriter(w, start)
	return binary.Write(sw, binary.LittleEndian, data)
}

func newFile(path string) (*os.File, *blockListInterface, error) {
	//This function creates a new file and writes out the header and initial block list
	bli := new(blockListInterface)
	bli.BlockListInfos = make(map[int64]*blockListInfo)
	file, err := os.Create(path)
	if err != nil {
		return file, nil, err
	}

	bli.file = file
	header := fileHeaderData{0, 0}
	bli.fileheader = &header
	manager, written, err := bli.newBlockList(file, int64(binary.Size(header)), freeBlockSize)

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

	bli.Blocklists.PushBack(manager)

	header.Freeblock_start = int64(binary.Size(header))
	header.Data_start = header.Freeblock_start + written
	err = writeTo(file, 0, header)
	return file, bli, err
}

func readFile(path string) (*os.File, *blockListInterface, error) {
	bli := new(blockListInterface)
	bli.BlockListInfos = make(map[int64]*blockListInfo)
	file, err := os.Open(path)
	if err != nil {
		return file, nil, err
	}

	bli.file = file
	header := new(fileHeaderData)
	bli.fileheader = header
	if err = readFrom(file, 0, header); err != nil {
		return file, nil, err
	}
	blm, err := bli.readBlockList(file, header.Freeblock_start)
	bli.Blocklists.PushBack(blm)
	for blm.header.Next != 0 {
		blm, err = bli.readBlockList(file, blm.header.Next)
		if err != nil {
			return file, nil, err
		}
		bli.Blocklists.PushBack(blm)
	}

	return file, bli, err
}
<commit_msg>Change way that file is opened<commit_after>package gokvlite

import (
	"container/list"
	"encoding/binary"
	"errors"
	"io"
	"os"
)

const freeBlockSize = 1024

type fileHeaderData struct {
	Freeblock_start int64
	Data_start      int64
}

type blockListHeaderData struct {
	Next int64
	Size int64
}

type blockListArrayEntryData struct {
	//This is a single array entry
	Free  uint8
	Start int64
	Size  int64
}

type blockListInfo struct {
	//This is a wrapper for a single array entry
	Entry    *blockListArrayEntryData
	Location int64
}

func (info *blockListInfo) writeInfo(writer io.WriterAt) error {
	return writeTo(writer, info.Location, info.Entry)
}

func (info *blockListInfo) ReadData(reader io.ReaderAt) (*[]byte, error) {
	//Reads the data from the location specific in the info/entry
	if info.Entry.Free > 0 {
		return nil, errors.New("filemanager: ReadData: Info is free, unable to read")
	}
	data := make([]byte, info.Entry.Size)
	_, err := reader.ReadAt(data, info.Entry.Start)
	return &data, err
}

func (info *blockListInfo) WriteData(writer io.WriterAt, data []byte) error {
	//Writes the data to the location specified by the info/entry
	if info.Entry.Free == 0 {
		return errors.New("Info is free, unable to write")
	}
	if int64(binary.Size(data)) != info.Entry.Size {
		return errors.New("Size of data is incorrect for size of info.")
	}
	_, err := writer.WriteAt(data, info.Entry.Start)
	return err
}

type blockListManager struct {
	headerStart int64
	header      *blockListHeaderData
}

type blockListInterface struct {
	/* This handles multiple block lists, gets a free block, etc
	 */

	fileheader     *fileHeaderData
	file           *os.File
	Blocklists     list.List
	Freeblocks     list.List
	Freeentries    list.List
	BlockListInfos map[int64]*blockListInfo
}

func (bli *blockListInterface) getFileEnd() (int64, error) {
	err := bli.file.Sync()
	if err != nil {
		return 0, err
	}

	fi, err := bli.file.Stat()
	if err != nil {
		return 0, err
	}

	end := fi.Size()
	return end, err
}

func (bli *blockListInterface) makeNewBlockList() error {
	end, err := bli.getFileEnd()
	if err != nil {
		return err
	}
	manager, _, err := bli.newBlockList(bli.file, end, freeBlockSize)
	if err != nil {
		return err
	}

	el := bli.Blocklists.Back()
	if el != nil {
		bl, ok := el.Value.(*blockListManager)
		if !ok {
			return errors.New("Invalid type for blocklistmanager in makeNewBlockList")
		}

		bl.header.Next = manager.headerStart
		err = writeTo(bli.file, bl.headerStart, bl.header)
		if err != nil {
			return err
		}
	}

	bli.Blocklists.PushBack(manager)
	return nil
}

func (bli *blockListInterface) getFreeEntry() (*blockListInfo, error) {
	//gets a free entry and removes it from the free entry list
	el := bli.Freeentries.Front()
	if el == nil {
		err := bli.makeNewBlockList()
		if err != nil {
			return nil, err
		}

		el = bli.Freeentries.Front()
		if el == nil {
			return nil, errors.New("Unable to get new element after making new blocklist")
		}
	}

	info, ok := el.Value.(*blockListInfo)
	if !ok {
		return info, errors.New("Invalid type in Freeentries list")
	}
	bli.Freeentries.Remove(el)

	return info, nil
}

func (bli *blockListInterface) Resize(info *blockListInfo, size int64) (*blockListInfo, *blockListInfo, error) {
	//resizes the info block, makes a new block, returns the old resized one and the new one
	startingsize := info.Entry.Size
	if startingsize == size {
		return nil, info, nil
	}

	if size > startingsize {
		//Just get a new entry and disregard the existing one
		newinfo, err := bli.GetFree(size)
		if err != nil {
			return nil, nil, err
		}
		err = bli.SetFree(info)
		if err != nil {
			return nil, nil, err
		}
		return nil, newinfo, nil
	}

	info.Entry.Size = size

	err := info.writeInfo(bli.file)
	if err != nil {
		return nil, nil, err
	}

	newinfo, err := bli.getFreeEntry()
	if err != nil {
		return nil, nil, err
	}

	newinfo.Entry.Size = startingsize - size
	newinfo.Entry.Free = 1
	newinfo.Entry.Start = info.Entry.Start + size
	err = newinfo.writeInfo(bli.file)
	if err != nil {
		return nil, nil, err
	}
	bli.Freeblocks.PushBack(newinfo)

	return info, newinfo, nil
}

func (bli *blockListInterface) SetFree(info *blockListInfo) error {
	//Sets a blocklist as free
	info.Entry.Free = 1
	err := info.writeInfo(bli.file)
	if err != nil {
		return err
	}
	bli.Freeblocks.PushBack(info)
	return nil
}

func (bli *blockListInterface) GetFree(size int64) (*blockListInfo, error) {
	//find a block large enough and resize it if it exists. Returns it marked as not free
	for e := bli.Freeblocks.Front(); e != nil; e = e.Next() {
		info, ok := e.Value.(*blockListInfo)
		if !ok {
			return nil, errors.New("Incorrect type in Freeblocks")
		}
		switch {
		case info.Entry.Size == size:
			info.Entry.Free = 0
			bli.Freeblocks.Remove(e)
			return info, nil

		case info.Entry.Size > size:
			_, _, err := bli.Resize(info, size)
			if err != nil {
				return nil, err
			}
			info.Entry.Free = 0
			bli.Freeblocks.Remove(e)
			return info, nil
		}
	}

	//append to end of file
	info, err := bli.getFreeEntry()
	if err != nil {
		return nil, err
	}

	end, err := bli.getFileEnd()
	if err != nil {
		return nil, err
	}

	info.Entry.Start = end
	info.Entry.Size = size
	info.Entry.Free = 0

	//Allocate the data storage
	data := make([]byte, size)
	bli.file.WriteAt(data, end)

	err = info.writeInfo(bli.file)
	return info, err
}

func (bli *blockListInterface) newBlockList(w io.WriterAt, start int64, size int64) (*blockListManager, int64, error) {
	//This creates a single unlinked blocklist
	//Size is in number of items, not bytes
	var written int64

	manager := new(blockListManager)

	header := new(blockListHeaderData)
	header.Size = size
	header.Next = 0
	err := writeTo(w, start, header)

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

	written += int64(binary.Size(header))

	manager.header = header
	manager.headerStart = start

	for i := int64(0); i < size; i++ {
		entry := new(blockListArrayEntryData)
		entry.Free = 1
		info := blockListInfo{entry, written + start}
		err = info.writeInfo(w)
		if err != nil {
			return manager, written, err
		}
		bli.Freeentries.PushBack(&info)
		bli.BlockListInfos[info.Location] = &info
		written += int64(binary.Size(entry))

	}
	return manager, written, err
}

func (bli *blockListInterface) readBlockList(reader io.ReaderAt, start int64) (*blockListManager, error) {
	var read int64
	header := new(blockListHeaderData)
	err := readFrom(reader, start, header)
	if err != nil {
		return nil, err
	}

	blm := new(blockListManager)
	blm.header = header
	blm.headerStart = start

	read += int64(binary.Size(header))

	for i := int64(0); i < blm.header.Size; i++ {
		data := new(blockListArrayEntryData)
		err := readFrom(reader, start+read, data)
		if err != nil {
			return blm, err
		}

		info := blockListInfo{data, start + read}
		if data.Size == 0 {
			bli.Freeentries.PushBack(&info)
		} else if data.Free > 0 {
			bli.Freeentries.PushBack(&info)
		}
		bli.BlockListInfos[info.Location] = &info

		read += int64(binary.Size(data))
	}
	return blm, nil
}

func readFrom(r io.ReaderAt, start int64, data interface{}) error {
	//This function reads any of the structs above and returns it
	sr := io.NewSectionReader(r, start, int64(binary.Size(data)))
	return binary.Read(sr, binary.LittleEndian, data)
}

func writeTo(w io.WriterAt, start int64, data interface{}) error {
	//This function writes any of the structs above and returns it
	sw := newSectionWriter(w, start)
	return binary.Write(sw, binary.LittleEndian, data)
}

func newFile(path string) (*os.File, *blockListInterface, error) {
	//This function creates a new file and writes out the header and initial block list
	bli := new(blockListInterface)
	bli.BlockListInfos = make(map[int64]*blockListInfo)
	file, err := os.OpenFile(path, os.O_RDWR | os.O_CREATE, os.FileMode(0666))
	if err != nil {
		return file, nil, err
	}

	bli.file = file
	header := fileHeaderData{0, 0}
	bli.fileheader = &header
	manager, written, err := bli.newBlockList(file, int64(binary.Size(header)), freeBlockSize)

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

	bli.Blocklists.PushBack(manager)

	header.Freeblock_start = int64(binary.Size(header))
	header.Data_start = header.Freeblock_start + written
	err = writeTo(file, 0, header)
	return file, bli, err
}

func readFile(path string) (*os.File, *blockListInterface, error) {
	bli := new(blockListInterface)
	bli.BlockListInfos = make(map[int64]*blockListInfo)
	file, err := os.OpenFile(path, os.O_RDWR, os.FileMode(0666))
	if err != nil {
		return file, nil, err
	}

	bli.file = file
	header := new(fileHeaderData)
	bli.fileheader = header
	if err = readFrom(file, 0, header); err != nil {
		return file, nil, err
	}
	blm, err := bli.readBlockList(file, header.Freeblock_start)
	bli.Blocklists.PushBack(blm)
	for blm.header.Next != 0 {
		blm, err = bli.readBlockList(file, blm.header.Next)
		if err != nil {
			return file, nil, err
		}
		bli.Blocklists.PushBack(blm)
	}

	return file, bli, err
}
<|endoftext|>


<commit_before>package hoverfly

import (
	"os"
	"sync"
)

// Configuration - initial structure of configuration
type Configuration struct {
	AdminPort    string
	ProxyPort    string
	Mode         string
	Destination  string
	Middleware   string
	DatabaseName string
	Verbose      bool
	Development  bool

	mu sync.Mutex
}

// SetMode - provides safe way to set new mode
func (c *Configuration) SetMode(mode string) {
	c.mu.Lock()
	c.Mode = mode
	c.mu.Unlock()
}

// GetMode - provides safe way to get current mode
func (c *Configuration) GetMode() (mode string) {
	c.mu.Lock()
	mode = c.Mode
	c.mu.Unlock()
	return
}

// DefaultPort - default proxy port
const DefaultPort = "8500"

// DefaultAdminPort - default admin interface port
const DefaultAdminPort = "8888"

// DefaultDatabaseName - default database name that will be created
// or used by Hoverfly
const DefaultDatabaseName = "requests.db"

// InitSettings gets and returns initial configuration from env
// variables or sets defaults
func InitSettings() *Configuration {

	var appConfig Configuration
	// getting default admin interface port
	if os.Getenv("AdminPort") != "" {
		appConfig.AdminPort = os.Getenv("AdminPort")
	} else {
		appConfig.AdminPort = DefaultAdminPort
	}

	// getting proxy port
	if os.Getenv("ProxyPort") != "" {
		appConfig.ProxyPort = os.Getenv("ProxyPort")
	} else {
		appConfig.ProxyPort = DefaultPort
	}

	databaseName := os.Getenv("HoverflyDB")
	if databaseName == "" {
		databaseName = DefaultDatabaseName
	}
	appConfig.DatabaseName = databaseName

	// middleware configuration
	appConfig.Middleware = os.Getenv("HoverflyMiddleware")

	return &appConfig
}
<commit_msg>settings updated<commit_after>package hoverfly

import (
	"os"
	"strconv"
	"sync"

	log "github.com/Sirupsen/logrus"
)

// Configuration - initial structure of configuration
type Configuration struct {
	AdminPort          string
	ProxyPort          string
	Mode               string
	Destination        string
	Middleware         string
	DatabaseName       string
	Verbose            bool
	Development        bool
	SecretKey          []byte
	JWTExpirationDelta int

	mu sync.Mutex
}

// SetMode - provides safe way to set new mode
func (c *Configuration) SetMode(mode string) {
	c.mu.Lock()
	c.Mode = mode
	c.mu.Unlock()
}

// GetMode - provides safe way to get current mode
func (c *Configuration) GetMode() (mode string) {
	c.mu.Lock()
	mode = c.Mode
	c.mu.Unlock()
	return
}

// DefaultPort - default proxy port
const DefaultPort = "8500"

// DefaultAdminPort - default admin interface port
const DefaultAdminPort = "8888"

// DefaultDatabaseName - default database name that will be created
// or used by Hoverfly
const DefaultDatabaseName = "requests.db"

const DefaultJWTExpirationDelta = 72

// InitSettings gets and returns initial configuration from env
// variables or sets defaults
func InitSettings() *Configuration {

	var appConfig Configuration
	// getting default admin interface port
	if os.Getenv("AdminPort") != "" {
		appConfig.AdminPort = os.Getenv("AdminPort")
	} else {
		appConfig.AdminPort = DefaultAdminPort
	}

	// getting proxy port
	if os.Getenv("ProxyPort") != "" {
		appConfig.ProxyPort = os.Getenv("ProxyPort")
	} else {
		appConfig.ProxyPort = DefaultPort
	}

	databaseName := os.Getenv("HoverflyDB")
	if databaseName == "" {
		databaseName = DefaultDatabaseName
	}
	appConfig.DatabaseName = databaseName

	if os.Getenv("HoverflySecret") != "" {
		appConfig.SecretKey = []byte(os.Getenv("HoverflySecret"))
	} else {
		appConfig.SecretKey = GetRandomName(10)
	}

	if os.Getenv("HoverflyTokenExpiration") != "" {

		exp, err := strconv.Atoi(os.Getenv("HoverflyTokenExpiration"))
		if err != nil {
			log.WithFields(log.Fields{
				"error":                   err.Error(),
				"HoverflyTokenExpiration": os.Getenv("HoverflyTokenExpiration"),
			}).Error("failed to get token exipration delta, using default value")
			exp = DefaultJWTExpirationDelta
		}
		appConfig.JWTExpirationDelta = exp

	} else {
		appConfig.JWTExpirationDelta = DefaultJWTExpirationDelta
	}

	// middleware configuration
	appConfig.Middleware = os.Getenv("HoverflyMiddleware")

	return &appConfig
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
)

const (
	VERSION = "0.9.0"
)

type AppSettings struct {
	verbose bool
	stats   bool

	splitOutput bool

	inputDummy  MultiOption
	outputDummy MultiOption

	inputTCP       MultiOption
	outputTCP      MultiOption
	outputTCPStats bool

	inputFile  MultiOption
	outputFile MultiOption

	inputRAW MultiOption

	inputHTTP                   MultiOption
	outputHTTP                  MultiOption
	outputHTTPHeaders           HTTPHeaders
	outputHTTPMethods           HTTPMethods
	outputHTTPUrlRegexp         HTTPUrlRegexp
	outputHTTPUrlRewrite        UrlRewriteMap
	outputHTTPHeaderFilters     HTTPHeaderFilters
	outputHTTPHeaderHashFilters HTTPHeaderHashFilters
	outputHTTPElasticSearch     string
	outputHTTPWorkers           int
	outputHTTPStats             bool
}

var Settings AppSettings = AppSettings{}

func usage() {
	fmt.Printf("Gor is a simple http traffic replication tool written in Go. Its main goal is to replay traffic from production servers to staging and dev environments.\nProject page: https://github.com/buger/gor\nAuthor: <Leonid Bugaev> leonsbox@gmail.com\nCurrent Version: %s\n\n", VERSION)
	flag.PrintDefaults()
	os.Exit(2)
}

func init() {
	flag.Usage = usage

	flag.BoolVar(&Settings.verbose, "verbose", false, "Turn on verbose/debug output")
	flag.BoolVar(&Settings.stats, "stats", false, "Turn on queue stats output")

	flag.BoolVar(&Settings.splitOutput, "split-output", false, "By default each output gets same traffic. If set to `true` it splits traffic equally among all outputs.")

	flag.Var(&Settings.inputDummy, "input-dummy", "Used for testing outputs. Emits 'Get /' request every 1s")
	flag.Var(&Settings.outputDummy, "output-dummy", "Used for testing inputs. Just prints data coming from inputs.")

	flag.Var(&Settings.inputTCP, "input-tcp", "Used for internal communication between Gor instances. Example: \n\t# Receive requests from other Gor instances on 28020 port, and redirect output to staging\n\tgor --input-tcp :28020 --output-http staging.com")
	flag.Var(&Settings.outputTCP, "output-tcp", "Used for internal communication between Gor instances. Example: \n\t# Listen for requests on 80 port and forward them to other Gor instance on 28020 port\n\tgor --input-raw :80 --output-tcp replay.local:28020")
	flag.BoolVar(&Settings.outputTCPStats, "output-tcp-stats", false, "Report TCP output queue stats to console every 5 seconds.")

	flag.Var(&Settings.inputFile, "input-file", "Read requests from file: \n\tgor --input-file ./requests.gor --output-http staging.com")
	flag.Var(&Settings.outputFile, "output-file", "Write incoming requests to file: \n\tgor --input-raw :80 --output-file ./requests.gor")

	flag.Var(&Settings.inputRAW, "input-raw", "Capture traffic from given port (use RAW sockets and require *sudo* access):\n\t# Capture traffic from 8080 port\n\tgor --input-raw :8080 --output-http staging.com")

	flag.Var(&Settings.inputHTTP, "input-http", "Read requests from HTTP, should be explicitly sent from your application:\n\t# Listen for http on 9000\n\tgor --input-http :9000 --output-http staging.com")

	flag.Var(&Settings.outputHTTP, "output-http", "Forwards incoming requests to given http address.\n\t# Redirect all incoming requests to staging.com address \n\tgor --input-raw :80 --output-http http://staging.com")
	flag.Var(&Settings.outputHTTPHeaders, "output-http-header", "Inject additional headers to http reqest:\n\tgor --input-raw :8080 --output-http staging.com --output-http-header 'User-Agent: Gor'")
	flag.Var(&Settings.outputHTTPMethods, "output-http-method", "Whitelist of HTTP methods to replay. Anything else will be dropped:\n\tgor --input-raw :8080 --output-http staging.com --output-http-method GET --output-http-method OPTIONS")
	flag.Var(&Settings.outputHTTPUrlRegexp, "output-http-url-regexp", "A regexp to match requests against. Anything else will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --output-http-url-regexp ^www.")
	flag.Var(&Settings.outputHTTPHeaderFilters, "output-http-header-filter", "A regexp to match a specific header against. Requests with non-matching headers will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --output-http-header-filter api-version:^v1")
	flag.Var(&Settings.outputHTTPHeaderHashFilters, "output-http-header-hash-filter", "Takes a fraction of requests, consistently taking or rejecting a request based on the FNV32-1A hash of a specific header. The fraction must have a denominator that is a power of two:\n\t gor --input-raw :8080 --output-http staging.com --output-http-header-hash-filter user-id:1/4")
	flag.IntVar(&Settings.outputHTTPWorkers, "output-http-workers", -1, "Gor uses dynamic worker scaling by default.  Enter a number to run a set number of workers.")
	flag.BoolVar(&Settings.outputHTTPStats, "output-http-stats", false, "Report http output queue stats to console every 5 seconds.")

	flag.StringVar(&Settings.outputHTTPElasticSearch, "output-http-elasticsearch", "", "Send request and response stats to ElasticSearch:\n\tgor --input-raw :8080 --output-http staging.com --output-http-elasticsearch 'es_host:api_port/index_name'")
	flag.Var(&Settings.outputHTTPUrlRewrite, "output-http-rewrite-url", "Rewrite the requst url based on a mapping:\n\tgor --input-raw :8080 --output-http staging.com --output-http-rewrite-url /xml_test/interface.php:/api/service.do")
}

func Debug(args ...interface{}) {
	if Settings.verbose {
		log.Println(args...)
	}
}
<commit_msg>Update version<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
)

const (
	VERSION = "0.9.3"
)

type AppSettings struct {
	verbose bool
	stats   bool

	splitOutput bool

	inputDummy  MultiOption
	outputDummy MultiOption

	inputTCP       MultiOption
	outputTCP      MultiOption
	outputTCPStats bool

	inputFile  MultiOption
	outputFile MultiOption

	inputRAW MultiOption

	inputHTTP                   MultiOption
	outputHTTP                  MultiOption
	outputHTTPHeaders           HTTPHeaders
	outputHTTPMethods           HTTPMethods
	outputHTTPUrlRegexp         HTTPUrlRegexp
	outputHTTPUrlRewrite        UrlRewriteMap
	outputHTTPHeaderFilters     HTTPHeaderFilters
	outputHTTPHeaderHashFilters HTTPHeaderHashFilters
	outputHTTPElasticSearch     string
	outputHTTPWorkers           int
	outputHTTPStats             bool
}

var Settings AppSettings = AppSettings{}

func usage() {
	fmt.Printf("Gor is a simple http traffic replication tool written in Go. Its main goal is to replay traffic from production servers to staging and dev environments.\nProject page: https://github.com/buger/gor\nAuthor: <Leonid Bugaev> leonsbox@gmail.com\nCurrent Version: %s\n\n", VERSION)
	flag.PrintDefaults()
	os.Exit(2)
}

func init() {
	flag.Usage = usage

	flag.BoolVar(&Settings.verbose, "verbose", false, "Turn on verbose/debug output")
	flag.BoolVar(&Settings.stats, "stats", false, "Turn on queue stats output")

	flag.BoolVar(&Settings.splitOutput, "split-output", false, "By default each output gets same traffic. If set to `true` it splits traffic equally among all outputs.")

	flag.Var(&Settings.inputDummy, "input-dummy", "Used for testing outputs. Emits 'Get /' request every 1s")
	flag.Var(&Settings.outputDummy, "output-dummy", "Used for testing inputs. Just prints data coming from inputs.")

	flag.Var(&Settings.inputTCP, "input-tcp", "Used for internal communication between Gor instances. Example: \n\t# Receive requests from other Gor instances on 28020 port, and redirect output to staging\n\tgor --input-tcp :28020 --output-http staging.com")
	flag.Var(&Settings.outputTCP, "output-tcp", "Used for internal communication between Gor instances. Example: \n\t# Listen for requests on 80 port and forward them to other Gor instance on 28020 port\n\tgor --input-raw :80 --output-tcp replay.local:28020")
	flag.BoolVar(&Settings.outputTCPStats, "output-tcp-stats", false, "Report TCP output queue stats to console every 5 seconds.")

	flag.Var(&Settings.inputFile, "input-file", "Read requests from file: \n\tgor --input-file ./requests.gor --output-http staging.com")
	flag.Var(&Settings.outputFile, "output-file", "Write incoming requests to file: \n\tgor --input-raw :80 --output-file ./requests.gor")

	flag.Var(&Settings.inputRAW, "input-raw", "Capture traffic from given port (use RAW sockets and require *sudo* access):\n\t# Capture traffic from 8080 port\n\tgor --input-raw :8080 --output-http staging.com")

	flag.Var(&Settings.inputHTTP, "input-http", "Read requests from HTTP, should be explicitly sent from your application:\n\t# Listen for http on 9000\n\tgor --input-http :9000 --output-http staging.com")

	flag.Var(&Settings.outputHTTP, "output-http", "Forwards incoming requests to given http address.\n\t# Redirect all incoming requests to staging.com address \n\tgor --input-raw :80 --output-http http://staging.com")
	flag.Var(&Settings.outputHTTPHeaders, "output-http-header", "Inject additional headers to http reqest:\n\tgor --input-raw :8080 --output-http staging.com --output-http-header 'User-Agent: Gor'")
	flag.Var(&Settings.outputHTTPMethods, "output-http-method", "Whitelist of HTTP methods to replay. Anything else will be dropped:\n\tgor --input-raw :8080 --output-http staging.com --output-http-method GET --output-http-method OPTIONS")
	flag.Var(&Settings.outputHTTPUrlRegexp, "output-http-url-regexp", "A regexp to match requests against. Anything else will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --output-http-url-regexp ^www.")
	flag.Var(&Settings.outputHTTPHeaderFilters, "output-http-header-filter", "A regexp to match a specific header against. Requests with non-matching headers will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --output-http-header-filter api-version:^v1")
	flag.Var(&Settings.outputHTTPHeaderHashFilters, "output-http-header-hash-filter", "Takes a fraction of requests, consistently taking or rejecting a request based on the FNV32-1A hash of a specific header. The fraction must have a denominator that is a power of two:\n\t gor --input-raw :8080 --output-http staging.com --output-http-header-hash-filter user-id:1/4")
	flag.IntVar(&Settings.outputHTTPWorkers, "output-http-workers", -1, "Gor uses dynamic worker scaling by default.  Enter a number to run a set number of workers.")
	flag.BoolVar(&Settings.outputHTTPStats, "output-http-stats", false, "Report http output queue stats to console every 5 seconds.")

	flag.StringVar(&Settings.outputHTTPElasticSearch, "output-http-elasticsearch", "", "Send request and response stats to ElasticSearch:\n\tgor --input-raw :8080 --output-http staging.com --output-http-elasticsearch 'es_host:api_port/index_name'")
	flag.Var(&Settings.outputHTTPUrlRewrite, "output-http-rewrite-url", "Rewrite the requst url based on a mapping:\n\tgor --input-raw :8080 --output-http staging.com --output-http-rewrite-url /xml_test/interface.php:/api/service.do")
}

func Debug(args ...interface{}) {
	if Settings.verbose {
		log.Println(args...)
	}
}
<|endoftext|>