<commit_before>package elog

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"runtime"
	"sync"
	"time"
)

const (
	lf          byte = 0x0A // 换行
	space       byte = 0x20 // 空格
	coreFormat       = "[%s] %v"
	timeLayout       = "2006-01-02 15:04:05.999"
	callerDepth      = 3
)

var (
	bufPool = &sync.Pool{
		New: func() interface{} {
			return &bytes.Buffer{}
		},
	}

	msgPool = &sync.Pool{
		New: func() interface{} {
			return &logMessage{}
		},
	}

	errLineNo = errors.New("ELog: get lineno encounter a error.")
)

func (e *ELog) baseLog(lvl LogLevel, msg string) {
	if e.cfg.LogLevel > lvl {
		return // 屏蔽打印
	}
	logMsg := msgPool.Get().(*logMessage)

	logMsg.msg = msg
	logMsg.lvl = lvl
	line, err := getLineNo()
	if err != nil {
		fmt.Fprintln(os.Stderr, "Get line number encounter a error.")
	}
	logMsg.lineNo = line

	e.waitWrite.Wait()
	e.logChan <- logMsg
}

func (e *ELog) Debug(format string, params ...interface{}) {
	e.baseLog(DebugLvl, fmt.Sprintf(format, params...))
}

func (e *ELog) Info(format string, params ...interface{}) {
	e.baseLog(InfoLvl, fmt.Sprintf(format, params...))
}

func (e *ELog) Warn(format string, params ...interface{}) {
	e.baseLog(WarnLvl, fmt.Sprintf(format, params...))
}

func (e *ELog) Error(format string, params ...interface{}) {
	e.baseLog(ErrorLvl, fmt.Sprintf(format, params...))
}

// Panic由调用者处理
func (e *ELog) Panic(format string, params ...interface{}) {
	e.baseLog(PanicLvl, fmt.Sprintf(format, params...))
}

// 这里会退出程序,调用os.Exit()
func (e *ELog) Fatal(format string, params ...interface{}) {
	e.baseLog(FatalLvl, fmt.Sprintf(format, params...))
}

// core func
func (e *ELog) log(logMsg *logMessage) {
	var buffer *bytes.Buffer
	buffer = bufPool.Get().(*bytes.Buffer)
	buffer.Reset() // 不能保证buffer是否被GC
	defer bufPool.Put(buffer)

	e.logger.Lock()
	defer e.logger.Unlock()

	// 前缀
	buffer.WriteString(e.logPrefix(logMsg.lvl))
	buffer.WriteByte(space)

	// 行号
	if e.cfg.ShowLineNumber {
		buffer.WriteString(logMsg.lineNo)
		buffer.WriteByte(space)
	}

	// 日志信息
	buffer.WriteString(logMsg.msg)
	buffer.WriteByte(lf)

	// 写入文件
	e.logger.f.Write(buffer.Bytes())

	// 写入标注输出
	if e.cfg.EnabledStdout {
		e.stdout.Write(buffer.Bytes())
	}
}

func (e *ELog) logPrefix(lvl LogLevel) string {
	now := time.Now()
	prefix := ""

	if len(e.cfg.TimeLayout) <= 0 {
		prefix = fmt.Sprintf(coreFormat, lvl.String(), now.Format(timeLayout))
	} else {
		prefix = fmt.Sprintf(coreFormat, lvl.String(), now.Format(e.cfg.TimeLayout))
	}

	return prefix
}

// 失败返回空字符串
func getLineNo() (string, error) {
	_, filePath, lineNo, ok := runtime.Caller(callerDepth)
	if !ok {
		return "", errLineNo
	}

	return fmt.Sprintf("%s:%d.", filePath, lineNo), nil
}
<commit_msg>avoid panic main channel.<commit_after>package elog

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"os"
	"runtime"
	"sync"
	"time"
)

const (
	lf          byte = 0x0A // 换行
	space       byte = 0x20 // 空格
	coreFormat       = "[%s] %v"
	timeLayout       = "2006-01-02 15:04:05.999"
	callerDepth      = 3
)

var (
	bufPool = &sync.Pool{
		New: func() interface{} {
			return &bytes.Buffer{}
		},
	}

	msgPool = &sync.Pool{
		New: func() interface{} {
			return &logMessage{}
		},
	}

	errLineNo = errors.New("ELog: get lineno encounter a error.")
)

func (e *ELog) baseLog(lvl LogLevel, msg string) {
	if e.cfg.LogLevel > lvl {
		return // 屏蔽打印
	}
	logMsg := msgPool.Get().(*logMessage)

	logMsg.msg = msg
	logMsg.lvl = lvl
	line, err := getLineNo()
	if err != nil {
		fmt.Fprintln(os.Stderr, "Get line number encounter a error.")
	}
	logMsg.lineNo = line

	e.waitWrite.Wait()
	e.logChan <- logMsg
}

func (e *ELog) Debug(format string, params ...interface{}) {
	e.baseLog(DebugLvl, fmt.Sprintf(format, params...))
}

func (e *ELog) Info(format string, params ...interface{}) {
	e.baseLog(InfoLvl, fmt.Sprintf(format, params...))
}

func (e *ELog) Warn(format string, params ...interface{}) {
	e.baseLog(WarnLvl, fmt.Sprintf(format, params...))
}

func (e *ELog) Error(format string, params ...interface{}) {
	e.baseLog(ErrorLvl, fmt.Sprintf(format, params...))
}

// Panic由调用者处理
func (e *ELog) Panic(format string, params ...interface{}) {
	e.baseLog(PanicLvl, fmt.Sprintf(format, params...))
}

// 这里会退出程序,调用os.Exit()
func (e *ELog) Fatal(format string, params ...interface{}) {
	e.baseLog(FatalLvl, fmt.Sprintf(format, params...))
}

// core func
func (e *ELog) log(logMsg *logMessage) {
	defer func() {
		if err := recover(); err != nil {
			// avoid panic main channel.
			log.Printf("Recover from a error.Error(%v)\n", err)
		}
	}()

	var buffer *bytes.Buffer
	buffer = bufPool.Get().(*bytes.Buffer)
	buffer.Reset() // 不能保证buffer是否被GC
	defer bufPool.Put(buffer)

	e.logger.Lock()
	defer e.logger.Unlock()

	// 前缀
	buffer.WriteString(e.logPrefix(logMsg.lvl))
	buffer.WriteByte(space)

	// 行号
	if e.cfg.ShowLineNumber {
		buffer.WriteString(logMsg.lineNo)
		buffer.WriteByte(space)
	}

	// 日志信息
	buffer.WriteString(logMsg.msg)
	buffer.WriteByte(lf)

	// 写入文件
	e.logger.f.Write(buffer.Bytes())

	// 写入标注输出
	if e.cfg.EnabledStdout {
		e.stdout.Write(buffer.Bytes())
	}
}

func (e *ELog) logPrefix(lvl LogLevel) string {
	now := time.Now()
	prefix := ""

	if len(e.cfg.TimeLayout) <= 0 {
		prefix = fmt.Sprintf(coreFormat, lvl.String(), now.Format(timeLayout))
	} else {
		prefix = fmt.Sprintf(coreFormat, lvl.String(), now.Format(e.cfg.TimeLayout))
	}

	return prefix
}

// 失败返回空字符串
func getLineNo() (string, error) {
	_, filePath, lineNo, ok := runtime.Caller(callerDepth)
	if !ok {
		return "", errLineNo
	}

	return fmt.Sprintf("%s:%d.", filePath, lineNo), nil
}
<|endoftext|>


<commit_before>package storage

import (
	"encoding/binary"
	"path/filepath"
	"time"

	"github.com/boltdb/bolt"

	"github.com/anacrolix/torrent/metainfo"
)

var (
	value = []byte{}
)

type boltPieceCompletion struct {
	db *bolt.DB
}

func NewBoltPieceCompletion(dir string) (ret PieceCompletion, err error) {
	p := filepath.Join(dir, ".torrent.bolt.db")
	db, err := bolt.Open(p, 0660, &bolt.Options{
		Timeout: time.Second,
	})
	if err != nil {
		return
	}
	ret = &boltPieceCompletion{db}
	return
}

func (me *boltPieceCompletion) Get(pk metainfo.PieceKey) (ret bool, err error) {
	err = me.db.View(func(tx *bolt.Tx) error {
		c := tx.Bucket(completed)
		if c == nil {
			return nil
		}
		ih := c.Bucket(pk.InfoHash[:])
		if ih == nil {
			return nil
		}
		var key [4]byte
		binary.BigEndian.PutUint32(key[:], uint32(pk.Index))
		ret = ih.Get(key[:]) != nil
		return nil
	})
	return
}

func (me *boltPieceCompletion) Set(pk metainfo.PieceKey, b bool) error {
	return me.db.Update(func(tx *bolt.Tx) error {
		c, err := tx.CreateBucketIfNotExists(completed)
		if err != nil {
			return err
		}
		ih, err := c.CreateBucketIfNotExists(pk.InfoHash[:])
		if err != nil {
			return err
		}
		var key [4]byte
		binary.BigEndian.PutUint32(key[:], uint32(pk.Index))
		if b {
			return ih.Put(key[:], value)
		} else {
			return ih.Delete(key[:])
		}
	})
}

func (me *boltPieceCompletion) Close() error {
	return me.db.Close()
}
<commit_msg>Make bolt completion DB directory if necessary<commit_after>package storage

import (
	"encoding/binary"
	"os"
	"path/filepath"
	"time"

	"github.com/boltdb/bolt"

	"github.com/anacrolix/torrent/metainfo"
)

var (
	value = []byte{}
)

type boltPieceCompletion struct {
	db *bolt.DB
}

func NewBoltPieceCompletion(dir string) (ret PieceCompletion, err error) {
	os.MkdirAll(dir, 0770)
	p := filepath.Join(dir, ".torrent.bolt.db")
	db, err := bolt.Open(p, 0660, &bolt.Options{
		Timeout: time.Second,
	})
	if err != nil {
		return
	}
	ret = &boltPieceCompletion{db}
	return
}

func (me *boltPieceCompletion) Get(pk metainfo.PieceKey) (ret bool, err error) {
	err = me.db.View(func(tx *bolt.Tx) error {
		c := tx.Bucket(completed)
		if c == nil {
			return nil
		}
		ih := c.Bucket(pk.InfoHash[:])
		if ih == nil {
			return nil
		}
		var key [4]byte
		binary.BigEndian.PutUint32(key[:], uint32(pk.Index))
		ret = ih.Get(key[:]) != nil
		return nil
	})
	return
}

func (me *boltPieceCompletion) Set(pk metainfo.PieceKey, b bool) error {
	return me.db.Update(func(tx *bolt.Tx) error {
		c, err := tx.CreateBucketIfNotExists(completed)
		if err != nil {
			return err
		}
		ih, err := c.CreateBucketIfNotExists(pk.InfoHash[:])
		if err != nil {
			return err
		}
		var key [4]byte
		binary.BigEndian.PutUint32(key[:], uint32(pk.Index))
		if b {
			return ih.Put(key[:], value)
		} else {
			return ih.Delete(key[:])
		}
	})
}

func (me *boltPieceCompletion) Close() error {
	return me.db.Close()
}
<|endoftext|>


<commit_before>package common

import (
	"regexp"
	"strings"

	"github.com/webx-top/com"
	"github.com/webx-top/db"
	"github.com/webx-top/db/lib/factory"
	"github.com/webx-top/db/mysql"
)

func SQLLineParser(exec func(string) error, lastSQLs ...*string) func(string) error {
	var sqlStr string
	return func(line string) error {
		if strings.HasPrefix(line, `--`) {
			return nil
		}
		line = strings.TrimRight(line, "\r ")
		if strings.HasPrefix(line, `/*`) && strings.HasSuffix(line, `*/;`) {
			return nil
		}
		sqlStr += line
		if strings.HasSuffix(line, `;`) {
			defer func() {
				sqlStr = ``
			}()
			//println(sqlStr)
			if sqlStr == `;` {
				return nil
			}
			return exec(sqlStr)
		}
		sqlStr += "\n"
		return nil
	}
}

func ParseSQL(sqlFile string, isFile bool, installer func(string) error) (err error) {
	installFunction := SQLLineParser(installer)
	if isFile {
		return com.SeekFileLines(sqlFile, installFunction)
	}
	sqlContent := sqlFile
	for _, line := range strings.Split(sqlContent, "\n") {
		err = installFunction(line)
		if err != nil {
			return err
		}
	}
	return err
}

// ReplacePrefix 替换前缀数据
func ReplacePrefix(m factory.Model, field string, oldPrefix string, newPrefix string) error {
	oldPrefix = com.AddSlashes(oldPrefix, '_', '%')
	value := db.Raw("REPLACE(`"+field+"`, ?, ?)", oldPrefix, newPrefix)
	return m.SetField(nil, field, value, field, db.Like(oldPrefix+`%`))
}

var (
	sqlCharsetRegexp     = regexp.MustCompile(`(?i) (CHARACTER SET |CHARSET=)utf8mb4 `)
	sqlCollateRegexp     = regexp.MustCompile(`(?i) (COLLATE[= ])utf8mb4_general_ci`)
	sqlCreateTableRegexp = regexp.MustCompile(`(?i)^CREATE TABLE `)
	mysqlNetworkRegexp   = regexp.MustCompile(`^[/]{2,}`)
)

// ReplaceCharset 替换DDL语句中的字符集
func ReplaceCharset(sqlStr string, charset string, checkCreateDDL ...bool) string {
	if charset == `utf8mb4` {
		return sqlStr
	}
	if len(checkCreateDDL) > 0 && checkCreateDDL[0] {
		if !sqlCreateTableRegexp.MatchString(sqlStr) {
			return sqlStr
		}
	}
	sqlStr = sqlCharsetRegexp.ReplaceAllString(sqlStr, ` ${1}`+charset+` `)
	sqlStr = sqlCollateRegexp.ReplaceAllString(sqlStr, ` ${1}`+charset+`_general_ci`)
	return sqlStr
}

func ParseMysqlConnectionURL(settings *mysql.ConnectionURL) {
	if strings.HasPrefix(settings.Host, `unix:`) {
		settings.Socket = strings.TrimPrefix(settings.Host, `unix:`)
		settings.Socket = mysqlNetworkRegexp.ReplaceAllString(settings.Socket, `/`)
		settings.Host = ``
	}
}
<commit_msg>update<commit_after>package common

import (
	"regexp"
	"strings"

	"github.com/webx-top/com"
	"github.com/webx-top/db"
	"github.com/webx-top/db/lib/factory"
	"github.com/webx-top/db/mysql"
)

func SQLLineParser(exec func(string) error) func(string) error {
	var sqlStr string
	return func(line string) error {
		if strings.HasPrefix(line, `--`) {
			return nil
		}
		line = strings.TrimRight(line, "\r ")
		if strings.HasPrefix(line, `/*`) && strings.HasSuffix(line, `*/;`) {
			return nil
		}
		sqlStr += line
		if strings.HasSuffix(line, `;`) {
			defer func() {
				sqlStr = ``
			}()
			//println(sqlStr)
			if sqlStr == `;` {
				return nil
			}
			return exec(sqlStr)
		}
		sqlStr += "\n"
		return nil
	}
}

func ParseSQL(sqlFile string, isFile bool, installer func(string) error) (err error) {
	installFunction := SQLLineParser(installer)
	if isFile {
		return com.SeekFileLines(sqlFile, installFunction)
	}
	sqlContent := sqlFile
	for _, line := range strings.Split(sqlContent, "\n") {
		err = installFunction(line)
		if err != nil {
			return err
		}
	}
	return err
}

// ReplacePrefix 替换前缀数据
func ReplacePrefix(m factory.Model, field string, oldPrefix string, newPrefix string) error {
	oldPrefix = com.AddSlashes(oldPrefix, '_', '%')
	value := db.Raw("REPLACE(`"+field+"`, ?, ?)", oldPrefix, newPrefix)
	return m.SetField(nil, field, value, field, db.Like(oldPrefix+`%`))
}

var (
	sqlCharsetRegexp     = regexp.MustCompile(`(?i) (CHARACTER SET |CHARSET=)utf8mb4 `)
	sqlCollateRegexp     = regexp.MustCompile(`(?i) (COLLATE[= ])utf8mb4_general_ci`)
	sqlCreateTableRegexp = regexp.MustCompile(`(?i)^CREATE TABLE `)
	mysqlNetworkRegexp   = regexp.MustCompile(`^[/]{2,}`)
)

// ReplaceCharset 替换DDL语句中的字符集
func ReplaceCharset(sqlStr string, charset string, checkCreateDDL ...bool) string {
	if charset == `utf8mb4` {
		return sqlStr
	}
	if len(checkCreateDDL) > 0 && checkCreateDDL[0] {
		if !sqlCreateTableRegexp.MatchString(sqlStr) {
			return sqlStr
		}
	}
	sqlStr = sqlCharsetRegexp.ReplaceAllString(sqlStr, ` ${1}`+charset+` `)
	sqlStr = sqlCollateRegexp.ReplaceAllString(sqlStr, ` ${1}`+charset+`_general_ci`)
	return sqlStr
}

func ParseMysqlConnectionURL(settings *mysql.ConnectionURL) {
	if strings.HasPrefix(settings.Host, `unix:`) {
		settings.Socket = strings.TrimPrefix(settings.Host, `unix:`)
		settings.Socket = mysqlNetworkRegexp.ReplaceAllString(settings.Socket, `/`)
		settings.Host = ``
	}
}
<|endoftext|>


<commit_before>// +build integration

package riak

import (
	"net"
	"testing"
	"time"
)

func TestCreateNodeWithOptionsAndStart(t *testing.T) {
	o := &testListenerOpts{
		test: t,
		host: "127.0.0.1",
		port: 13340,
	}
	tl := newTestListener(o)
	tl.start()
	defer tl.stop()

	count := uint16(16)
	opts := &NodeOptions{
		RemoteAddress:       tl.addr,
		MinConnections:      count,
		MaxConnections:      count,
		IdleTimeout:         thirtySeconds,
		ConnectTimeout:      thirtySeconds,
		RequestTimeout:      thirtySeconds,
		HealthCheckInterval: time.Millisecond * 500,
		HealthCheckBuilder:  &PingCommandBuilder{},
	}
	node, err := NewNode(opts)
	if err != nil {
		t.Error(err.Error())
	}
	if node == nil {
		t.Fatal("expected non-nil node")
	}
	if node.addr.Port != int(tl.port) {
		t.Errorf("expected port %d, got: %d", tl.port, node.addr.Port)
	}
	if node.addr.Zone != "" {
		t.Errorf("expected empty zone, got: %s", string(node.addr.Zone))
	}
	if expected, actual := opts.MinConnections, node.cm.minConnections; expected != actual {
		t.Errorf("expected %v, got: %v", expected, actual)
	}
	if expected, actual := opts.MaxConnections, node.cm.maxConnections; expected != actual {
		t.Errorf("expected %v, got: %v", expected, actual)
	}
	if expected, actual := opts.IdleTimeout, node.cm.idleTimeout; expected != actual {
		t.Errorf("expected %v, got: %v", expected, actual)
	}
	if err := node.start(); err != nil {
		t.Error(err)
	}
	var f = func(v interface{}) (bool, bool) {
		conn := v.(*connection)
		if conn == nil {
			t.Error("got unexpected nil value")
			return true, false
		}
		if expected, actual := int(tl.port), conn.addr.Port; expected != actual {
			t.Errorf("expected %d, got: %d", expected, actual)
		}
		if conn.addr.Zone != "" {
			t.Errorf("expected empty zone, got: %s", string(conn.addr.Zone))
		}
		if conn.healthCheck != nil {
			t.Error("expected nil conn.healthCheck")
		}
		if expected, actual := conn.connectTimeout, opts.ConnectTimeout; expected != actual {
			t.Errorf("expected %v, got: %v", expected, actual)
		}
		if expected, actual := conn.requestTimeout, opts.RequestTimeout; expected != actual {
			t.Errorf("expected %v, got: %v", expected, actual)
		}
		return false, true
	}
	if err := node.cm.q.iterate(f); err != nil {
		t.Error(err)
	}
	if err := node.stop(); err != nil {
		t.Error(err)
	}
}

func TestRecoverViaDefaultPingHealthCheck(t *testing.T) {
	connects := 0
	var onConn = func(c net.Conn) bool {
		connects++
		if connects == 1 {
			c.Close()
		} else {
			readWritePingResp(t, c, true)
		}
		return true
	}
	o := &testListenerOpts{
		test:   t,
		host:   "127.0.0.1",
		port:   13337,
		onConn: onConn,
	}
	tl := newTestListener(o)
	tl.start()
	defer tl.stop()

	doneChan := make(chan struct{})
	stateChan := make(chan state)

	go func() {
		opts := &NodeOptions{
			RemoteAddress:       tl.addr,
			MinConnections:      0,
			HealthCheckInterval: 50 * time.Millisecond,
		}
		node, err := NewNode(opts)
		if err != nil {
			t.Error(err)
		}

		origSetStateFunc := node.setStateFunc
		node.setStateFunc = func(sd *stateData, st state) {
			origSetStateFunc(&node.stateData, st)
			logDebug("[TestRecoverViaDefaultPingHealthCheck]", "sending state '%v' down stateChan", st)
			stateChan <- st
		}

		node.start()
		ping := &PingCommand{}
		executed, err := node.execute(ping)
		if executed == false {
			t.Fatal("expected ping to be executed")
		}
		if err == nil {
			t.Fatal("expected non-nil error")
		}
		logDebug("[TestRecoverViaDefaultPingHealthCheck]", "waiting to stop node")
		<-doneChan
		logDebug("[TestRecoverViaDefaultPingHealthCheck]", "stopping node")
		node.stop()
	}()

	checkStatesFunc := func(states []state) {
		for i := 0; i < len(states); i++ {
			nodeState := <-stateChan
			if expected, actual := states[i], nodeState; expected != actual {
				t.Errorf("expected %v, got %v", expected, actual)
			} else {
				logDebug("[TestRecoverViaDefaultPingHealthCheck]", "saw state %d", nodeState)
			}
		}
	}

	expectedStates := []state{
		nodeRunning, nodeHealthChecking, nodeRunning,
	}

	checkStatesFunc(expectedStates)

	close(doneChan)

	expectedStates = []state{
		nodeShuttingDown, nodeShutdown,
	}

	checkStatesFunc(expectedStates)

	close(stateChan)
}

func TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck(t *testing.T) {
	o := &testListenerOpts{
		test: t,
		host: "127.0.0.1",
		port: 13338,
	}
	tl := newTestListener(o)
	defer tl.stop()

	stateChan := make(chan state)
	recoveredChan := make(chan struct{})

	var node *Node
	opts := &NodeOptions{
		ConnectTimeout: 125 * time.Millisecond,
		RemoteAddress:  tl.addr,
	}
	var err error
	node, err = NewNode(opts)
	if err != nil {
		t.Fatal(err)
	}
	origSetStateFunc := node.setStateFunc

	go func() {
		node.setStateFunc = func(sd *stateData, st state) {
			origSetStateFunc(&node.stateData, st)
			logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "sending state '%v' down stateChan", st)
			stateChan <- st
		}
		node.start()

		pc := &PingCommand{}
		node.execute(pc)
		for {
			select {
			case <-recoveredChan:
				break
			case <-time.After(time.Second):
				logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "waiting for recovery...")
				pc := &PingCommand{}
				node.execute(pc)
			}
		}
	}()

	go func() {
		listenerStarted := false
		nodeIsRunningCount := 0
		for {
			if nodeState, ok := <-stateChan; ok {
				logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "nodeState: '%v'", nodeState)
				if node.isCurrentState(nodeRunning) {
					nodeIsRunningCount++
				}
				if nodeIsRunningCount == 2 {
					// This is the second time node has entered nodeRunning state, so it must have recovered via the health check
					logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "SUCCESS node recovered via health check")
					close(recoveredChan)
					break
				}
				if !listenerStarted && node.isCurrentState(nodeHealthChecking) {
					logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "STARTING LISTENER")
					tl.start()
				}
			} else {
				t.Error("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck] stateChan closed before recovering via health check")
				break
			}
		}
	}()

	select {
	case <-recoveredChan:
		logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "recovered")
		node.setStateFunc = origSetStateFunc
		node.stop()
		close(stateChan)
	case <-time.After(5 * time.Second):
		t.Error("test timed out")
	}
}
<commit_msg>Give test a bit more time<commit_after>// +build integration

package riak

import (
	"net"
	"testing"
	"time"
)

func TestCreateNodeWithOptionsAndStart(t *testing.T) {
	o := &testListenerOpts{
		test: t,
		host: "127.0.0.1",
		port: 13340,
	}
	tl := newTestListener(o)
	tl.start()
	defer tl.stop()

	count := uint16(16)
	opts := &NodeOptions{
		RemoteAddress:       tl.addr,
		MinConnections:      count,
		MaxConnections:      count,
		IdleTimeout:         thirtySeconds,
		ConnectTimeout:      thirtySeconds,
		RequestTimeout:      thirtySeconds,
		HealthCheckInterval: time.Millisecond * 500,
		HealthCheckBuilder:  &PingCommandBuilder{},
	}
	node, err := NewNode(opts)
	if err != nil {
		t.Error(err.Error())
	}
	if node == nil {
		t.Fatal("expected non-nil node")
	}
	if node.addr.Port != int(tl.port) {
		t.Errorf("expected port %d, got: %d", tl.port, node.addr.Port)
	}
	if node.addr.Zone != "" {
		t.Errorf("expected empty zone, got: %s", string(node.addr.Zone))
	}
	if expected, actual := opts.MinConnections, node.cm.minConnections; expected != actual {
		t.Errorf("expected %v, got: %v", expected, actual)
	}
	if expected, actual := opts.MaxConnections, node.cm.maxConnections; expected != actual {
		t.Errorf("expected %v, got: %v", expected, actual)
	}
	if expected, actual := opts.IdleTimeout, node.cm.idleTimeout; expected != actual {
		t.Errorf("expected %v, got: %v", expected, actual)
	}
	if err := node.start(); err != nil {
		t.Error(err)
	}
	var f = func(v interface{}) (bool, bool) {
		conn := v.(*connection)
		if conn == nil {
			t.Error("got unexpected nil value")
			return true, false
		}
		if expected, actual := int(tl.port), conn.addr.Port; expected != actual {
			t.Errorf("expected %d, got: %d", expected, actual)
		}
		if conn.addr.Zone != "" {
			t.Errorf("expected empty zone, got: %s", string(conn.addr.Zone))
		}
		if conn.healthCheck != nil {
			t.Error("expected nil conn.healthCheck")
		}
		if expected, actual := conn.connectTimeout, opts.ConnectTimeout; expected != actual {
			t.Errorf("expected %v, got: %v", expected, actual)
		}
		if expected, actual := conn.requestTimeout, opts.RequestTimeout; expected != actual {
			t.Errorf("expected %v, got: %v", expected, actual)
		}
		return false, true
	}
	if err := node.cm.q.iterate(f); err != nil {
		t.Error(err)
	}
	if err := node.stop(); err != nil {
		t.Error(err)
	}
}

func TestRecoverViaDefaultPingHealthCheck(t *testing.T) {
	connects := 0
	var onConn = func(c net.Conn) bool {
		connects++
		if connects == 1 {
			c.Close()
		} else {
			readWritePingResp(t, c, true)
		}
		return true
	}
	o := &testListenerOpts{
		test:   t,
		host:   "127.0.0.1",
		port:   13337,
		onConn: onConn,
	}
	tl := newTestListener(o)
	tl.start()
	defer tl.stop()

	doneChan := make(chan struct{})
	stateChan := make(chan state)

	go func() {
		opts := &NodeOptions{
			RemoteAddress:       tl.addr,
			MinConnections:      0,
			HealthCheckInterval: 50 * time.Millisecond,
		}
		node, err := NewNode(opts)
		if err != nil {
			t.Error(err)
		}

		origSetStateFunc := node.setStateFunc
		node.setStateFunc = func(sd *stateData, st state) {
			origSetStateFunc(&node.stateData, st)
			logDebug("[TestRecoverViaDefaultPingHealthCheck]", "sending state '%v' down stateChan", st)
			stateChan <- st
		}

		node.start()
		ping := &PingCommand{}
		executed, err := node.execute(ping)
		if executed == false {
			t.Fatal("expected ping to be executed")
		}
		if err == nil {
			t.Fatal("expected non-nil error")
		}
		logDebug("[TestRecoverViaDefaultPingHealthCheck]", "waiting to stop node")
		<-doneChan
		logDebug("[TestRecoverViaDefaultPingHealthCheck]", "stopping node")
		node.stop()
	}()

	checkStatesFunc := func(states []state) {
		for i := 0; i < len(states); i++ {
			nodeState := <-stateChan
			if expected, actual := states[i], nodeState; expected != actual {
				t.Errorf("expected %v, got %v", expected, actual)
			} else {
				logDebug("[TestRecoverViaDefaultPingHealthCheck]", "saw state %d", nodeState)
			}
		}
	}

	expectedStates := []state{
		nodeRunning, nodeHealthChecking, nodeRunning,
	}

	checkStatesFunc(expectedStates)

	close(doneChan)

	expectedStates = []state{
		nodeShuttingDown, nodeShutdown,
	}

	checkStatesFunc(expectedStates)

	close(stateChan)
}

func TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck(t *testing.T) {
	o := &testListenerOpts{
		test: t,
		host: "127.0.0.1",
		port: 13338,
	}
	tl := newTestListener(o)
	defer tl.stop()

	stateChan := make(chan state)
	recoveredChan := make(chan struct{})

	var node *Node
	opts := &NodeOptions{
		ConnectTimeout: 125 * time.Millisecond,
		RemoteAddress:  tl.addr,
	}
	var err error
	node, err = NewNode(opts)
	if err != nil {
		t.Fatal(err)
	}
	origSetStateFunc := node.setStateFunc

	go func() {
		node.setStateFunc = func(sd *stateData, st state) {
			origSetStateFunc(&node.stateData, st)
			logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "sending state '%v' down stateChan", st)
			stateChan <- st
		}
		node.start()

		pc := &PingCommand{}
		node.execute(pc)
		for {
			select {
			case <-recoveredChan:
				break
			case <-time.After(time.Second):
				logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "waiting for recovery...")
				pc := &PingCommand{}
				node.execute(pc)
			}
		}
	}()

	go func() {
		listenerStarted := false
		nodeIsRunningCount := 0
		for {
			if nodeState, ok := <-stateChan; ok {
				logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "nodeState: '%v'", nodeState)
				if node.isCurrentState(nodeRunning) {
					nodeIsRunningCount++
				}
				if nodeIsRunningCount == 2 {
					// This is the second time node has entered nodeRunning state, so it must have recovered via the health check
					logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "SUCCESS node recovered via health check")
					close(recoveredChan)
					break
				}
				if !listenerStarted && node.isCurrentState(nodeHealthChecking) {
					logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "STARTING LISTENER")
					tl.start()
				}
			} else {
				t.Error("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck] stateChan closed before recovering via health check")
				break
			}
		}
	}()

	select {
	case <-recoveredChan:
		logDebug("[TestRecoverAfterConnectionComesUpViaDefaultPingHealthCheck]", "recovered")
		node.setStateFunc = origSetStateFunc
		node.stop()
		close(stateChan)
	case <-time.After(10 * time.Second):
		t.Error("test timed out")
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 Vector Creations Ltd
//
// 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 routing

import (
	"net/http"
	"time"

	"github.com/gorilla/mux"
	"github.com/matrix-org/dendrite/clientapi/auth/storage/accounts"
	"github.com/matrix-org/dendrite/clientapi/auth/storage/devices"
	"github.com/matrix-org/dendrite/clientapi/producers"
	"github.com/matrix-org/dendrite/common"
	"github.com/matrix-org/dendrite/common/config"
	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"
)

const (
	pathPrefixV2Keys       = "/_matrix/key/v2"
	pathPrefixV1Federation = "/_matrix/federation/v1"
)

// Setup registers HTTP handlers with the given ServeMux.
func Setup(
	apiMux *mux.Router,
	cfg config.Dendrite,
	query api.RoomserverQueryAPI,
	aliasAPI api.RoomserverAliasAPI,
	producer *producers.RoomserverProducer,
	keys gomatrixserverlib.KeyRing,
	federation *gomatrixserverlib.FederationClient,
	accountDB *accounts.Database,
	deviceDB *devices.Database,
) {
	v2keysmux := apiMux.PathPrefix(pathPrefixV2Keys).Subrouter()
	v1fedmux := apiMux.PathPrefix(pathPrefixV1Federation).Subrouter()

	localKeys := common.MakeExternalAPI("localkeys", func(req *http.Request) util.JSONResponse {
		return LocalKeys(cfg)
	})

	// Ignore the {keyID} argument as we only have a single server key so we always
	// return that key.
	// Even if we had more than one server key, we would probably still ignore the
	// {keyID} argument and always return a response containing all of the keys.
	v2keysmux.Handle("/server/{keyID}", localKeys).Methods(http.MethodGet)
	v2keysmux.Handle("/server/", localKeys).Methods(http.MethodGet)

	v1fedmux.Handle("/send/{txnID}/", common.MakeFedAPI(
		"federation_send", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return Send(
				httpReq, request, gomatrixserverlib.TransactionID(vars["txnID"]),
				cfg, query, producer, keys, federation,
			)
		},
	)).Methods(http.MethodPut, http.MethodOptions)

	v1fedmux.Handle("/invite/{roomID}/{eventID}", common.MakeFedAPI(
		"federation_invite", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return Invite(
				httpReq, request, vars["roomID"], vars["eventID"],
				cfg, producer, keys,
			)
		},
	)).Methods(http.MethodPut, http.MethodOptions)

	v1fedmux.Handle("/3pid/onbind", common.MakeExternalAPI("3pid_onbind",
		func(req *http.Request) util.JSONResponse {
			return CreateInvitesFrom3PIDInvites(req, query, cfg, producer, federation, accountDB)
		},
	)).Methods(http.MethodPost, http.MethodOptions)

	v1fedmux.Handle("/exchange_third_party_invite/{roomID}", common.MakeFedAPI(
		"exchange_third_party_invite", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return ExchangeThirdPartyInvite(
				httpReq, request, vars["roomID"], query, cfg, federation, producer,
			)
		},
	)).Methods(http.MethodPut, http.MethodOptions)

	v1fedmux.Handle("/event/{eventID}", common.MakeFedAPI(
		"federation_get_event", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return GetEvent(
				httpReq.Context(), request, query, vars["eventID"],
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/state/{roomID}", common.MakeFedAPI(
		"federation_get_event_auth", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return GetState(
				httpReq.Context(), request, cfg, query, time.Now(),
				keys, vars["roomID"],
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/state_ids/{roomID}", common.MakeFedAPI(
		"federation_get_event_auth", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return GetStateIDs(
				httpReq.Context(), request, cfg, query, time.Now(),
				keys, vars["roomID"],
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/query/directory/", common.MakeFedAPI(
		"federation_query_room_alias", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			return RoomAliasToID(
				httpReq, federation, cfg, aliasAPI,
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/query/profile", common.MakeFedAPI(
		"federation_query_profile", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			return GetProfile(
				httpReq, accountDB, cfg,
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/query/user_devices/{userID}", common.MakeFedAPI(
		"federation_query_user_devices", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return GetUserDevices(
				httpReq, deviceDB, vars["userID"],
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/make_join/{roomID}/{userID}", common.MakeFedAPI(
		"federation_make_join", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			roomID := vars["roomID"]
			userID := vars["userID"]
			return MakeJoin(
				httpReq, request, cfg, query, roomID, userID,
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/send_join/{roomID}/{userID}", common.MakeFedAPI(
		"federation_send_join", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			roomID := vars["roomID"]
			userID := vars["userID"]
			return SendJoin(
				httpReq.Context(), httpReq, request, cfg, query, producer, keys, roomID, userID,
			)
		},
	)).Methods(http.MethodPut)

	v1fedmux.Handle("/make_leave/{roomID}/{userID}", common.MakeFedAPI(
		"federation_make_leave", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			roomID := vars["roomID"]
			userID := vars["userID"]
			return MakeLeave(
				httpReq, request, cfg, query, roomID, userID,
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/send_leave/{roomID}/{userID}", common.MakeFedAPI(
		"federation_send_leave", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			roomID := vars["roomID"]
			userID := vars["userID"]
			return SendLeave(
				httpReq, request, cfg, producer, keys, roomID, userID,
			)
		},
	)).Methods(http.MethodPut)

	v1fedmux.Handle("/version", common.MakeExternalAPI(
		"federation_version",
		func(httpReq *http.Request) util.JSONResponse {
			return Version()
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("get_missing_events/{roomID}", common.MakeFedAPI(
		"federation_get_missing_events", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return GetMissingEvents(httpReq, request, query, vars["roomID"])
		},
	)).Methods(http.MethodGet)
}
<commit_msg>Correct user/devices path (#557)<commit_after>// Copyright 2017 Vector Creations Ltd
//
// 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 routing

import (
	"net/http"
	"time"

	"github.com/gorilla/mux"
	"github.com/matrix-org/dendrite/clientapi/auth/storage/accounts"
	"github.com/matrix-org/dendrite/clientapi/auth/storage/devices"
	"github.com/matrix-org/dendrite/clientapi/producers"
	"github.com/matrix-org/dendrite/common"
	"github.com/matrix-org/dendrite/common/config"
	"github.com/matrix-org/dendrite/roomserver/api"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/util"
)

const (
	pathPrefixV2Keys       = "/_matrix/key/v2"
	pathPrefixV1Federation = "/_matrix/federation/v1"
)

// Setup registers HTTP handlers with the given ServeMux.
func Setup(
	apiMux *mux.Router,
	cfg config.Dendrite,
	query api.RoomserverQueryAPI,
	aliasAPI api.RoomserverAliasAPI,
	producer *producers.RoomserverProducer,
	keys gomatrixserverlib.KeyRing,
	federation *gomatrixserverlib.FederationClient,
	accountDB *accounts.Database,
	deviceDB *devices.Database,
) {
	v2keysmux := apiMux.PathPrefix(pathPrefixV2Keys).Subrouter()
	v1fedmux := apiMux.PathPrefix(pathPrefixV1Federation).Subrouter()

	localKeys := common.MakeExternalAPI("localkeys", func(req *http.Request) util.JSONResponse {
		return LocalKeys(cfg)
	})

	// Ignore the {keyID} argument as we only have a single server key so we always
	// return that key.
	// Even if we had more than one server key, we would probably still ignore the
	// {keyID} argument and always return a response containing all of the keys.
	v2keysmux.Handle("/server/{keyID}", localKeys).Methods(http.MethodGet)
	v2keysmux.Handle("/server/", localKeys).Methods(http.MethodGet)

	v1fedmux.Handle("/send/{txnID}/", common.MakeFedAPI(
		"federation_send", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return Send(
				httpReq, request, gomatrixserverlib.TransactionID(vars["txnID"]),
				cfg, query, producer, keys, federation,
			)
		},
	)).Methods(http.MethodPut, http.MethodOptions)

	v1fedmux.Handle("/invite/{roomID}/{eventID}", common.MakeFedAPI(
		"federation_invite", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return Invite(
				httpReq, request, vars["roomID"], vars["eventID"],
				cfg, producer, keys,
			)
		},
	)).Methods(http.MethodPut, http.MethodOptions)

	v1fedmux.Handle("/3pid/onbind", common.MakeExternalAPI("3pid_onbind",
		func(req *http.Request) util.JSONResponse {
			return CreateInvitesFrom3PIDInvites(req, query, cfg, producer, federation, accountDB)
		},
	)).Methods(http.MethodPost, http.MethodOptions)

	v1fedmux.Handle("/exchange_third_party_invite/{roomID}", common.MakeFedAPI(
		"exchange_third_party_invite", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return ExchangeThirdPartyInvite(
				httpReq, request, vars["roomID"], query, cfg, federation, producer,
			)
		},
	)).Methods(http.MethodPut, http.MethodOptions)

	v1fedmux.Handle("/event/{eventID}", common.MakeFedAPI(
		"federation_get_event", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return GetEvent(
				httpReq.Context(), request, query, vars["eventID"],
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/state/{roomID}", common.MakeFedAPI(
		"federation_get_event_auth", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return GetState(
				httpReq.Context(), request, cfg, query, time.Now(),
				keys, vars["roomID"],
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/state_ids/{roomID}", common.MakeFedAPI(
		"federation_get_event_auth", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return GetStateIDs(
				httpReq.Context(), request, cfg, query, time.Now(),
				keys, vars["roomID"],
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/query/directory/", common.MakeFedAPI(
		"federation_query_room_alias", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			return RoomAliasToID(
				httpReq, federation, cfg, aliasAPI,
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/query/profile", common.MakeFedAPI(
		"federation_query_profile", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			return GetProfile(
				httpReq, accountDB, cfg,
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/user/devices/{userID}", common.MakeFedAPI(
		"federation_user_devices", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return GetUserDevices(
				httpReq, deviceDB, vars["userID"],
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/make_join/{roomID}/{userID}", common.MakeFedAPI(
		"federation_make_join", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			roomID := vars["roomID"]
			userID := vars["userID"]
			return MakeJoin(
				httpReq, request, cfg, query, roomID, userID,
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/send_join/{roomID}/{userID}", common.MakeFedAPI(
		"federation_send_join", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			roomID := vars["roomID"]
			userID := vars["userID"]
			return SendJoin(
				httpReq.Context(), httpReq, request, cfg, query, producer, keys, roomID, userID,
			)
		},
	)).Methods(http.MethodPut)

	v1fedmux.Handle("/make_leave/{roomID}/{userID}", common.MakeFedAPI(
		"federation_make_leave", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			roomID := vars["roomID"]
			userID := vars["userID"]
			return MakeLeave(
				httpReq, request, cfg, query, roomID, userID,
			)
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("/send_leave/{roomID}/{userID}", common.MakeFedAPI(
		"federation_send_leave", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			roomID := vars["roomID"]
			userID := vars["userID"]
			return SendLeave(
				httpReq, request, cfg, producer, keys, roomID, userID,
			)
		},
	)).Methods(http.MethodPut)

	v1fedmux.Handle("/version", common.MakeExternalAPI(
		"federation_version",
		func(httpReq *http.Request) util.JSONResponse {
			return Version()
		},
	)).Methods(http.MethodGet)

	v1fedmux.Handle("get_missing_events/{roomID}", common.MakeFedAPI(
		"federation_get_missing_events", cfg.Matrix.ServerName, keys,
		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest) util.JSONResponse {
			vars := mux.Vars(httpReq)
			return GetMissingEvents(httpReq, request, query, vars["roomID"])
		},
	)).Methods(http.MethodGet)
}
<|endoftext|>


<commit_before>package v1

import (
	"database/sql"
	"strconv"
	"time"

	"git.zxq.co/ripple/rippleapi/common"
	"git.zxq.co/ripple/rippleapi/limit"
)

type rankRequestsStatusResponse struct {
	common.ResponseBase
	QueueSize       int        `json:"queue_size"`
	MaxPerUser      int        `json:"max_per_user"`
	Submitted       int        `json:"submitted"`
	SubmittedByUser *int       `json:"submitted_by_user,omitempty"`
	CanSubmit       *bool      `json:"can_submit,omitempty"`
	NextExpiration  *time.Time `json:"next_expiration"`
}

// BeatmapRankRequestsStatusGET gets the current status for beatmap ranking requests.
func BeatmapRankRequestsStatusGET(md common.MethodData) common.CodeMessager {
	c := common.GetConf()
	rows, err := md.DB.Query("SELECT userid, time FROM rank_requests WHERE time > ? ORDER BY id ASC LIMIT "+strconv.Itoa(c.RankQueueSize), time.Now().Add(-time.Hour*24).Unix())
	if err != nil {
		md.Err(err)
		return Err500
	}
	var r rankRequestsStatusResponse
	// if it's not auth-free access and we have got ReadConfidential, we can
	// know if this user can submit beatmaps or not.
	hasConfid := md.ID() != 0 && md.User.TokenPrivileges&common.PrivilegeReadConfidential > 0
	if hasConfid {
		r.SubmittedByUser = new(int)
	}
	isFirst := true
	for rows.Next() {
		var (
			user      int
			timestamp common.UnixTimestamp
		)
		err := rows.Scan(&user, &timestamp)
		if err != nil {
			md.Err(err)
			continue
		}
		// if the user submitted this rank request, increase the number of
		// rank requests submitted by this user
		if user == md.ID() && r.SubmittedByUser != nil {
			(*r.SubmittedByUser)++
		}
		// also, if this is the first result, it means it will be the next to
		// expire.
		if isFirst {
			x := time.Time(timestamp)
			r.NextExpiration = &x
			isFirst = false
		}
		r.Submitted++
	}
	r.QueueSize = c.RankQueueSize
	r.MaxPerUser = c.BeatmapRequestsPerUser
	if hasConfid {
		x := r.Submitted < r.QueueSize && *r.SubmittedByUser < r.MaxPerUser
		r.CanSubmit = &x
	}
	r.Code = 200
	return r
}

type submitRequestData struct {
	ID    int `json:"id"`
	SetID int `json:"set_id"`
}

// BeatmapRankRequestsSubmitPOST submits a new beatmap for ranking approval.
func BeatmapRankRequestsSubmitPOST(md common.MethodData) common.CodeMessager {
	var d submitRequestData
	err := md.RequestData.Unmarshal(&d)
	if err != nil {
		return ErrBadJSON
	}
	// check json data is present
	if d.ID == 0 && d.SetID == 0 {
		return ErrMissingField("id|set_id")
	}

	// you've been rate limited
	if !limit.NonBlockingRequest("rankrequest:u:"+strconv.Itoa(md.ID()), 5) {
		return common.SimpleResponse(429, "You may only try to request 5 beatmaps per minute.")
	}
	if !limit.NonBlockingRequest("rankrequest:ip:"+md.C.ClientIP(), 8) {
		return common.SimpleResponse(429, "You may only try to request 8 beatmaps per minute from the same IP.")
	}

	// find out from BeatmapRankRequestsStatusGET if we can submit beatmaps.
	statusRaw := BeatmapRankRequestsStatusGET(md)
	status, ok := statusRaw.(rankRequestsStatusResponse)
	if !ok {
		// if it's not a rankRequestsStatusResponse, it means it's an error
		return statusRaw
	}
	if !*status.CanSubmit {
		return common.SimpleResponse(403, "It's not possible to do a rank request at this time.")
	}

	w := common.
		Where("beatmap_id = ?", strconv.Itoa(d.ID)).Or().
		Where("beatmapset_id = ?", strconv.Itoa(d.SetID))

	var ranked int
	err = md.DB.QueryRow("SELECT ranked FROM beatmaps "+w.Clause+" LIMIT 1", w.Params...).Scan(&ranked)
	if ranked >= 2 {
		return common.SimpleResponse(406, "That beatmap is already ranked.")
	}

	switch err {
	case nil:
		// move on
	case sql.ErrNoRows:
		if d.SetID != 0 {
			md.R.Publish("lets:beatmap_updates:sets", strconv.Itoa(d.SetID))
		} else {
			md.R.Publish("lets:beatmap_updates:single", strconv.Itoa(d.ID))
		}
	default:
		md.Err(err)
		return Err500
	}

	// type and value of beatmap rank request
	t := "b"
	v := d.ID
	if d.SetID != 0 {
		t = "s"
		v = d.SetID
	}
	err = md.DB.QueryRow("SELECT 1 FROM rank_requests WHERE bid = ? AND type = ? AND time > ?",
		v, t, time.Now().Add(-time.Hour*24).Unix()).Scan(new(int))

	// error handling
	switch err {
	case sql.ErrNoRows:
		break
	case nil:
		// we're returning a success because if the request was already sent in the past 24
		// hours, it's as if the user submitted it.
		return BeatmapRankRequestsStatusGET(md)
	default:
		md.Err(err)
		return Err500
	}

	_, err = md.DB.Exec(
		"INSERT INTO rank_requests (userid, bid, type, time, blacklisted) VALUES (?, ?, ?, ?, 0)",
		md.ID(), v, t, time.Now().Unix())
	if err != nil {
		md.Err(err)
		return Err500
	}

	return BeatmapRankRequestsStatusGET(md)
}
<commit_msg>Use JSON instead of GLI ZINGONI<commit_after>package v1

import (
	"database/sql"
	"encoding/json"
	"strconv"
	"time"

	"git.zxq.co/ripple/rippleapi/common"
	"git.zxq.co/ripple/rippleapi/limit"
)

type rankRequestsStatusResponse struct {
	common.ResponseBase
	QueueSize       int        `json:"queue_size"`
	MaxPerUser      int        `json:"max_per_user"`
	Submitted       int        `json:"submitted"`
	SubmittedByUser *int       `json:"submitted_by_user,omitempty"`
	CanSubmit       *bool      `json:"can_submit,omitempty"`
	NextExpiration  *time.Time `json:"next_expiration"`
}

// BeatmapRankRequestsStatusGET gets the current status for beatmap ranking requests.
func BeatmapRankRequestsStatusGET(md common.MethodData) common.CodeMessager {
	c := common.GetConf()
	rows, err := md.DB.Query("SELECT userid, time FROM rank_requests WHERE time > ? ORDER BY id ASC LIMIT "+strconv.Itoa(c.RankQueueSize), time.Now().Add(-time.Hour*24).Unix())
	if err != nil {
		md.Err(err)
		return Err500
	}
	var r rankRequestsStatusResponse
	// if it's not auth-free access and we have got ReadConfidential, we can
	// know if this user can submit beatmaps or not.
	hasConfid := md.ID() != 0 && md.User.TokenPrivileges&common.PrivilegeReadConfidential > 0
	if hasConfid {
		r.SubmittedByUser = new(int)
	}
	isFirst := true
	for rows.Next() {
		var (
			user      int
			timestamp common.UnixTimestamp
		)
		err := rows.Scan(&user, &timestamp)
		if err != nil {
			md.Err(err)
			continue
		}
		// if the user submitted this rank request, increase the number of
		// rank requests submitted by this user
		if user == md.ID() && r.SubmittedByUser != nil {
			(*r.SubmittedByUser)++
		}
		// also, if this is the first result, it means it will be the next to
		// expire.
		if isFirst {
			x := time.Time(timestamp)
			r.NextExpiration = &x
			isFirst = false
		}
		r.Submitted++
	}
	r.QueueSize = c.RankQueueSize
	r.MaxPerUser = c.BeatmapRequestsPerUser
	if hasConfid {
		x := r.Submitted < r.QueueSize && *r.SubmittedByUser < r.MaxPerUser
		r.CanSubmit = &x
	}
	r.Code = 200
	return r
}

type submitRequestData struct {
	ID    int `json:"id"`
	SetID int `json:"set_id"`
}

// BeatmapRankRequestsSubmitPOST submits a new beatmap for ranking approval.
func BeatmapRankRequestsSubmitPOST(md common.MethodData) common.CodeMessager {
	var d submitRequestData
	err := md.RequestData.Unmarshal(&d)
	if err != nil {
		return ErrBadJSON
	}
	// check json data is present
	if d.ID == 0 && d.SetID == 0 {
		return ErrMissingField("id|set_id")
	}

	// you've been rate limited
	if !limit.NonBlockingRequest("rankrequest:u:"+strconv.Itoa(md.ID()), 5) {
		return common.SimpleResponse(429, "You may only try to request 5 beatmaps per minute.")
	}
	if !limit.NonBlockingRequest("rankrequest:ip:"+md.C.ClientIP(), 8) {
		return common.SimpleResponse(429, "You may only try to request 8 beatmaps per minute from the same IP.")
	}

	// find out from BeatmapRankRequestsStatusGET if we can submit beatmaps.
	statusRaw := BeatmapRankRequestsStatusGET(md)
	status, ok := statusRaw.(rankRequestsStatusResponse)
	if !ok {
		// if it's not a rankRequestsStatusResponse, it means it's an error
		return statusRaw
	}
	if !*status.CanSubmit {
		return common.SimpleResponse(403, "It's not possible to do a rank request at this time.")
	}

	w := common.
		Where("beatmap_id = ?", strconv.Itoa(d.ID)).Or().
		Where("beatmapset_id = ?", strconv.Itoa(d.SetID))

	var ranked int
	err = md.DB.QueryRow("SELECT ranked FROM beatmaps "+w.Clause+" LIMIT 1", w.Params...).Scan(&ranked)
	if ranked >= 2 {
		return common.SimpleResponse(406, "That beatmap is already ranked.")
	}

	switch err {
	case nil:
		// move on
	case sql.ErrNoRows:
		data, _ := json.Marshal(d)
		md.R.Publish("lets:beatmap_updates", string(data))
	default:
		md.Err(err)
		return Err500
	}

	// type and value of beatmap rank request
	t := "b"
	v := d.ID
	if d.SetID != 0 {
		t = "s"
		v = d.SetID
	}
	err = md.DB.QueryRow("SELECT 1 FROM rank_requests WHERE bid = ? AND type = ? AND time > ?",
		v, t, time.Now().Add(-time.Hour*24).Unix()).Scan(new(int))

	// error handling
	switch err {
	case sql.ErrNoRows:
		break
	case nil:
		// we're returning a success because if the request was already sent in the past 24
		// hours, it's as if the user submitted it.
		return BeatmapRankRequestsStatusGET(md)
	default:
		md.Err(err)
		return Err500
	}

	_, err = md.DB.Exec(
		"INSERT INTO rank_requests (userid, bid, type, time, blacklisted) VALUES (?, ?, ?, ?, 0)",
		md.ID(), v, t, time.Now().Unix())
	if err != nil {
		md.Err(err)
		return Err500
	}

	return BeatmapRankRequestsStatusGET(md)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"sort"

	"image"
	"image/draw"
	"image/png"
)

type Size struct {
	x, y int
}

func (a Size) Equal(b Size) bool {
	return a.x == b.x && a.y == b.y
}

func (a Size) Larger(b Size) bool {
	return b.x < a.x && b.y < a.y
}

func (a Size) Smaller(b Size) bool {
	return !a.Larger(b)
}

type sprite struct {
	name string
	img  image.Image
	rect image.Rectangle
	area int
	size Size
}

type ByArea []sprite

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

type Node struct {
	child [2]*Node
	rect  image.Rectangle
	img   *sprite
}

func (n *Node) size() Size {
	return Size{n.rect.Dx(), n.rect.Dy()}
}

func (n *Node) print() {
	fmt.Println(n)
	if n.child[0] != nil {
		n.child[0].print()
	}
	if n.child[1] != nil {
		n.child[1].print()
	}
}

func (n *Node) insert(img *sprite) bool {
	// there is already an image in this node
	if n.img != nil {
		fmt.Println("already contains an image")
		return false
	}

	// try to insert into either of the nodes children
	if n.child[0] != nil {
		fmt.Println("has a 0 child")
		in := n.child[0].insert(img)

		if in {
			return true
		} else {
			fmt.Println("has a 1 child")
			return n.child[1].insert(img)
		}
	}

	if n.rect.Dx() < img.size.x || n.rect.Dy() < img.size.y {
		fmt.Println("space too small")
		return false
	}

	if n.rect.Dx() == img.size.x && n.rect.Dy() == img.size.y {
		fmt.Println("prefect fit")
		n.img = img
		return true
	}

	if n.rect.Dx() >= img.size.x && n.rect.Dy() >= img.size.y {
		fmt.Println("Split")
		n.split(img)
	}

	return n.insert(img)

}

func (n *Node) split(img *sprite) {
	var tl0 image.Point
	var br0 image.Point

	var tl1 image.Point
	var br1 image.Point

	dx := n.size().x - img.size.x
	dy := n.size().y - img.size.y

	rc := n.rect

	tl0 = rc.Min
	br1 = rc.Max

	if dx > dy {
		fmt.Println("split on x")
		br0 = image.Point{rc.Min.X + img.size.x, rc.Dy()}
		tl1 = image.Point{rc.Min.X + img.size.x - 1, rc.Min.Y}
	} else {
		fmt.Println("split on y")
		br0 = image.Point{rc.Dx(), rc.Min.Y + img.size.y}
		tl1 = image.Point{rc.Min.X, rc.Min.Y + img.size.y - 1}
	}

	rect0 := image.Rectangle{tl0, br0}
	n.child[0] = &Node{rect: rect0}

	rect1 := image.Rectangle{tl1, br1}
	n.child[1] = &Node{rect: rect1}
}

func main() {
	flag.Parse()

	args := flag.Args()

	inputDir := args[0]

	files, _ := ioutil.ReadDir(inputDir)
	sprites := make([]sprite, len(files))

	totalX := 0
	totalY := 0

	for i := range sprites {
		s := readSprite(inputDir, files[i].Name())
		sprites[i] = s
		totalX += s.size.x
		totalY += s.size.y
	}

	// we want to place the largest sprite first
	sort.Sort(sort.Reverse(ByArea(sprites)))

	// the final image
	dst := image.NewRGBA(image.Rect(0, 0, 2048, 2048))

	n := Node{rect: image.Rect(0, 0, 1024, 1024)}

	for i := range sprites {
		s := &sprites[i]
		fmt.Printf("inserting %s\n", s.name)
		n.insert(s)
		draw.Draw(dst, s.rect, s.img, image.ZP, draw.Src)
	}

	n.print()

	writer, err := os.Create("test.png")
	err = png.Encode(writer, dst)
	if err != nil {
		log.Fatal(err)
	}
}

func readSprite(dir, name string) (s sprite) {
	path := path.Join(dir, name)
	reader, err := os.Open(path)
	if err != nil {
		log.Fatal(err)
	}
	defer reader.Close()

	img, err := png.Decode(reader)
	if err != nil {
		log.Fatal(err)
	}

	s.name = name
	s.img = img
	s.rect = img.Bounds()
	s.size = Size{s.rect.Dx(), s.rect.Dy()}
	s.area = s.size.x * s.size.y
	return
}
<commit_msg>Return fitting rectangle from insert<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path"
	"sort"

	"image"
	"image/draw"
	"image/png"
)

type Size struct {
	x, y int
}

func (a Size) Equal(b Size) bool {
	return a.x == b.x && a.y == b.y
}

func (a Size) Larger(b Size) bool {
	return b.x < a.x && b.y < a.y
}

func (a Size) Smaller(b Size) bool {
	return !a.Larger(b)
}

type sprite struct {
	name string
	img  image.Image
	rect image.Rectangle
	area int
	size Size
}

type ByArea []sprite

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

type Node struct {
	child [2]*Node
	rect  image.Rectangle
	img   *sprite
}

func (n *Node) size() Size {
	return Size{n.rect.Dx(), n.rect.Dy()}
}

func (n *Node) print() {
	fmt.Println(n)
	if n.child[0] != nil {
		n.child[0].print()
	}
	if n.child[1] != nil {
		n.child[1].print()
	}
}

func (n *Node) insert(img *sprite) (*image.Rectangle, error) {
	// there is already an image in this node
	if n.img != nil {
		fmt.Println("already contains an image")
		return nil, errors.New("already contains an image")
	}

	// try to insert into either of the nodes children
	if n.child[0] != nil {
		fmt.Println("has a 0 child")
		rc, err := n.child[0].insert(img)

		if err != nil {
			fmt.Println("has a 1 child")
			return n.child[1].insert(img)
		} else {
			return rc, nil
		}
	}

	if n.rect.Dx() < img.size.x || n.rect.Dy() < img.size.y {
		fmt.Println("space too small")
		return nil, errors.New("space too small")
	}

	if n.rect.Dx() == img.size.x && n.rect.Dy() == img.size.y {
		fmt.Println("prefect fit")
		n.img = img
		return &n.rect, nil
	}

	if n.rect.Dx() >= img.size.x && n.rect.Dy() >= img.size.y {
		fmt.Println("Split")
		n.split(img)
	}

	return n.insert(img)

}

func (n *Node) split(img *sprite) {
	var tl0 image.Point
	var br0 image.Point

	var tl1 image.Point
	var br1 image.Point

	dx := n.size().x - img.size.x
	dy := n.size().y - img.size.y

	rc := n.rect

	tl0 = rc.Min
	br1 = rc.Max

	if dx > dy {
		fmt.Println("split on x")
		br0 = image.Point{rc.Min.X + img.size.x, rc.Dy()}
		tl1 = image.Point{rc.Min.X + img.size.x - 1, rc.Min.Y}
	} else {
		fmt.Println("split on y")
		br0 = image.Point{rc.Dx(), rc.Min.Y + img.size.y}
		tl1 = image.Point{rc.Min.X, rc.Min.Y + img.size.y - 1}
	}

	rect0 := image.Rectangle{tl0, br0}
	n.child[0] = &Node{rect: rect0}

	rect1 := image.Rectangle{tl1, br1}
	n.child[1] = &Node{rect: rect1}
}

func main() {
	flag.Parse()

	args := flag.Args()

	inputDir := args[0]

	files, _ := ioutil.ReadDir(inputDir)
	sprites := make([]sprite, len(files))

	totalX := 0
	totalY := 0

	for i := range sprites {
		s := readSprite(inputDir, files[i].Name())
		sprites[i] = s
		totalX += s.size.x
		totalY += s.size.y
	}

	// we want to place the largest sprite first
	sort.Sort(sort.Reverse(ByArea(sprites)))

	// the final image
	dst := image.NewRGBA(image.Rect(0, 0, 2048, 2048))

	n := Node{rect: image.Rect(0, 0, 1024, 1024)}

	for i := range sprites {
		s := &sprites[i]
		fmt.Printf("inserting %s\n", s.name)
		n.insert(s)
		draw.Draw(dst, s.rect, s.img, image.ZP, draw.Src)
	}

	n.print()

	writer, err := os.Create("test.png")
	err = png.Encode(writer, dst)
	if err != nil {
		log.Fatal(err)
	}
}

func readSprite(dir, name string) (s sprite) {
	path := path.Join(dir, name)
	reader, err := os.Open(path)
	if err != nil {
		log.Fatal(err)
	}
	defer reader.Close()

	img, err := png.Decode(reader)
	if err != nil {
		log.Fatal(err)
	}

	s.name = name
	s.img = img
	s.rect = img.Bounds()
	s.size = Size{s.rect.Dx(), s.rect.Dy()}
	s.area = s.size.x * s.size.y
	return
}
<|endoftext|>


<commit_before>package handlers

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strconv"

	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/user"

	"github.com/icco/natnatnat/models"
	"github.com/pilu/traffic"
)

type ImportStruct struct {
	Id       int
	Datetime time.Time
	Text     string
	Title    string
}

func ImportTumbleHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	u := user.Current(c)
	if u == nil {
		url, _ := user.LoginURL(c, "/post/new")
		http.Redirect(w, r.Request, url, 302)
		return
	} else {
		log.Infof(c, "Logged in as: %s", u.String())
	}

	if u != nil && !user.IsAdmin(c) {
		http.Error(w, errors.New("Not a valid user.").Error(), 403)
		return
	} else {
		file, err := ioutil.ReadFile("tumbledata.json")
		if err != nil {
			e := fmt.Sprintf("File error: %v", err)
			http.Error(w, e, 500)
			return
		}

		var data []ImportStruct
		json.Unmarshal(file, &data)
		log.Debugf("Loaded: %v", data)

		w.WriteText("Finished.")
	}
}
<commit_msg>fix imports<commit_after>package handlers

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/user"

	"github.com/icco/natnatnat/models"
	"github.com/pilu/traffic"
)

type ImportStruct struct {
	Id       int
	Datetime time.Time
	Text     string
	Title    string
}

func ImportTumbleHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	u := user.Current(c)
	md := models.Markdown("test")
	if u == nil {
		url, _ := user.LoginURL(c, "/post/new")
		http.Redirect(w, r.Request, url, 302)
		return
	} else {
		log.Infof(c, "Logged in as: %s", u.String())
	}

	if u != nil && !user.IsAdmin(c) {
		http.Error(w, errors.New("Not a valid user.").Error(), 403)
		return
	} else {
		file, err := ioutil.ReadFile("tumbledata.json")
		if err != nil {
			e := fmt.Sprintf("File error: %v", err)
			http.Error(w, e, 500)
			return
		}

		var data []ImportStruct
		json.Unmarshal(file, &data)
		log.Debugf("Loaded: %v", data)

		w.WriteText("Finished.")
	}
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/lakeformation"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
	"github.com/hashicorp/terraform-plugin-sdk/helper/validation"
)

func resourceAwsLakeFormationResource() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsLakeFormationResourceRegister,
		Read:   resourceAwsLakeFormationResourceDescribe,
		Delete: resourceAwsLakeFormationResourceDeregister,

		Schema: map[string]*schema.Schema{
			"resource_arn": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validation.NoZeroValues,
			},
			"role_arn": {
				Type:         schema.TypeString,
				Optional:     true,
				Computed:     true,
				ForceNew:     true,
				ValidateFunc: validation.NoZeroValues,
			},
			"use_service_linked_role": {
				Type:     schema.TypeBool,
				Required: true,
				ForceNew: true,
			},
			"last_modified": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsLakeFormationResourceRegister(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	resourceArn := d.Get("resource_arn").(string)
	useServiceLinkedRole := d.Get("use_service_linked_role").(bool)

	input := &lakeformation.RegisterResourceInput{
		ResourceArn:          aws.String(resourceArn),
		UseServiceLinkedRole: aws.Bool(useServiceLinkedRole),
	}
	if v, ok := d.GetOk("role_arn"); ok {
		input.RoleArn = aws.String(v.(string))
	}

	_, err := conn.RegisterResource(input)
	if err != nil {
		return fmt.Errorf("Error registering LakeFormation Resource: %s", err)
	}

	d.SetId(fmt.Sprintf("lakeformation:resource:%s", resourceArn))

	return resourceAwsLakeFormationResourceDescribe(d, meta)
}

func resourceAwsLakeFormationResourceDescribe(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	resourceArn := d.Get("resource_arn").(string)

	input := &lakeformation.DescribeResourceInput{
		ResourceArn: aws.String(resourceArn),
	}

	out, err := conn.DescribeResource(input)
	if err != nil {
		return fmt.Errorf("Error reading LakeFormation Resource: %s", err)
	}

	d.Set("resource_arn", resourceArn)
	d.Set("role_arn", out.ResourceInfo.RoleArn)
	if out.ResourceInfo.LastModified != nil {
		d.Set("last_modified", out.ResourceInfo.LastModified.Format(time.RFC3339))
	}

	return nil
}

func resourceAwsLakeFormationResourceDeregister(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	resourceArn := d.Get("resource_arn").(string)

	input := &lakeformation.DeregisterResourceInput{
		ResourceArn: aws.String(resourceArn),
	}

	_, err := conn.DeregisterResource(input)
	if err != nil {
		return fmt.Errorf("Error deregistering LakeFormation Resource: %s", err)
	}

	return nil
}
<commit_msg>Better validations<commit_after>package aws

import (
	"fmt"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/lakeformation"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
)

func resourceAwsLakeFormationResource() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsLakeFormationResourceRegister,
		Read:   resourceAwsLakeFormationResourceDescribe,
		Delete: resourceAwsLakeFormationResourceDeregister,

		Schema: map[string]*schema.Schema{
			"resource_arn": {
				Type:         schema.TypeString,
				Required:     true,
				ForceNew:     true,
				ValidateFunc: validateArn,
			},
			"role_arn": {
				Type:         schema.TypeString,
				Optional:     true,
				Computed:     true,
				ForceNew:     true,
				ValidateFunc: validateArn,
			},
			"use_service_linked_role": {
				Type:     schema.TypeBool,
				Required: true,
				ForceNew: true,
			},
			"last_modified": {
				Type:     schema.TypeString,
				Computed: true,
			},
		},
	}
}

func resourceAwsLakeFormationResourceRegister(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	resourceArn := d.Get("resource_arn").(string)
	useServiceLinkedRole := d.Get("use_service_linked_role").(bool)

	input := &lakeformation.RegisterResourceInput{
		ResourceArn:          aws.String(resourceArn),
		UseServiceLinkedRole: aws.Bool(useServiceLinkedRole),
	}
	if v, ok := d.GetOk("role_arn"); ok {
		input.RoleArn = aws.String(v.(string))
	}

	_, err := conn.RegisterResource(input)
	if err != nil {
		return fmt.Errorf("Error registering LakeFormation Resource: %s", err)
	}

	d.SetId(fmt.Sprintf("lakeformation:resource:%s", resourceArn))

	return resourceAwsLakeFormationResourceDescribe(d, meta)
}

func resourceAwsLakeFormationResourceDescribe(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	resourceArn := d.Get("resource_arn").(string)

	input := &lakeformation.DescribeResourceInput{
		ResourceArn: aws.String(resourceArn),
	}

	out, err := conn.DescribeResource(input)
	if err != nil {
		return fmt.Errorf("Error reading LakeFormation Resource: %s", err)
	}

	d.Set("resource_arn", resourceArn)
	d.Set("role_arn", out.ResourceInfo.RoleArn)
	if out.ResourceInfo.LastModified != nil {
		d.Set("last_modified", out.ResourceInfo.LastModified.Format(time.RFC3339))
	}

	return nil
}

func resourceAwsLakeFormationResourceDeregister(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).lakeformationconn
	resourceArn := d.Get("resource_arn").(string)

	input := &lakeformation.DeregisterResourceInput{
		ResourceArn: aws.String(resourceArn),
	}

	_, err := conn.DeregisterResource(input)
	if err != nil {
		return fmt.Errorf("Error deregistering LakeFormation Resource: %s", err)
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package unicode

// Bit masks for each code point under U+0100, for fast lookup.
const (
	pC     = 1 << iota // a control character.
	pP                 // a punctuation character.
	pN                 // a numeral.
	pS                 // a symbolic character.
	pZ                 // a spacing character.
	pLu                // an upper-case letter.
	pLl                // a lower-case letter.
	pp                 // a printable character according to Go's definition.
	pg     = pp | pZ   // a graphical character according to the Unicode definition.
	pLo    = pLl | pLu // a letter that is neither upper nor lower case.
	pLmask = pLo
)

// GraphicRanges defines the set of graphic characters according to Unicode.
var GraphicRanges = []*RangeTable{
	L, M, N, P, S, Zs,
}

// PrintRanges defines the set of printable characters according to Go.
// ASCII space, U+0020, is handled separately.
var PrintRanges = []*RangeTable{
	L, M, N, P, S,
}

// IsGraphic reports whether the rune is defined as a Graphic by Unicode.
// Such characters include letters, marks, numbers, punctuation, symbols, and
// spaces, from categories L, M, N, P, S, Zs.
func IsGraphic(r rune) bool {
	// We convert to uint32 to avoid the extra test for negative,
	// and in the index we convert to uint8 to avoid the range check.
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pg != 0
	}
	return In(r, GraphicRanges...)
}

// IsPrint reports whether the rune is defined as printable by Go. Such
// characters include letters, marks, numbers, punctuation, symbols, and the
// ASCII space character, from categories L, M, N, P, S and the ASCII space
// character.  This categorization is the same as IsGraphic except that the
// only spacing character is ASCII space, U+0020.
func IsPrint(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pp != 0
	}
	return In(r, PrintRanges...)
}

// IsOneOf reports whether the rune is a member of one of the ranges.
// The function "In" provides a nicer signature and should be used in preference to IsOneOf.
func IsOneOf(ranges []*RangeTable, r rune) bool {
	for _, inside := range ranges {
		if Is(inside, r) {
			return true
		}
	}
	return false
}

// In reports whether the rune is a member of one of the ranges.
func In(r rune, ranges ...*RangeTable) bool {
	for _, inside := range ranges {
		if Is(inside, r) {
			return true
		}
	}
	return false
}

// IsControl reports whether the rune is a control character.
// The C (Other) Unicode category includes more code points
// such as surrogates; use Is(C, r) to test for them.
func IsControl(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pC != 0
	}
	// All control characters are < Latin1Max.
	return false
}

// IsLetter reports whether the rune is a letter (category L).
func IsLetter(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&(pLmask) != 0
	}
	return isExcludingLatin(Letter, r)
}

// IsMark reports whether the rune is a mark character (category M).
func IsMark(r rune) bool {
	// There are no mark characters in Latin-1.
	return isExcludingLatin(Mark, r)
}

// IsNumber reports whether the rune is a number (category N).
func IsNumber(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pN != 0
	}
	return isExcludingLatin(Number, r)
}

// IsPunct reports whether the rune is a Unicode punctuation character
// (category P).
func IsPunct(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pP != 0
	}
	return Is(Punct, r)
}

// IsSpace reports whether the rune is a space character as defined
// by Unicode's White Space property; in the Latin-1 space
// this is
//	'\t', '\n', '\v', '\f', '\r', ' ', U+0085 (NEL), U+00A0 (NBSP).
// Other definitions of spacing characters are set by category
// Z and property Pattern_White_Space.
func IsSpace(r rune) bool {
	// This property isn't the same as Z; special-case it.
	if uint32(r) <= MaxLatin1 {
		switch r {
		case '\t', '\n', '\v', '\f', '\r', ' ', 0x85, 0xA0:
			return true
		}
		return false
	}
	return isExcludingLatin(White_Space, r)
}

// IsSymbol reports whether the rune is a symbolic character.
func IsSymbol(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pS != 0
	}
	return isExcludingLatin(Symbol, r)
}
<commit_msg>unicode: Fixed an out of date comment (MaxLatin1, not Latin1Max).<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package unicode

// Bit masks for each code point under U+0100, for fast lookup.
const (
	pC     = 1 << iota // a control character.
	pP                 // a punctuation character.
	pN                 // a numeral.
	pS                 // a symbolic character.
	pZ                 // a spacing character.
	pLu                // an upper-case letter.
	pLl                // a lower-case letter.
	pp                 // a printable character according to Go's definition.
	pg     = pp | pZ   // a graphical character according to the Unicode definition.
	pLo    = pLl | pLu // a letter that is neither upper nor lower case.
	pLmask = pLo
)

// GraphicRanges defines the set of graphic characters according to Unicode.
var GraphicRanges = []*RangeTable{
	L, M, N, P, S, Zs,
}

// PrintRanges defines the set of printable characters according to Go.
// ASCII space, U+0020, is handled separately.
var PrintRanges = []*RangeTable{
	L, M, N, P, S,
}

// IsGraphic reports whether the rune is defined as a Graphic by Unicode.
// Such characters include letters, marks, numbers, punctuation, symbols, and
// spaces, from categories L, M, N, P, S, Zs.
func IsGraphic(r rune) bool {
	// We convert to uint32 to avoid the extra test for negative,
	// and in the index we convert to uint8 to avoid the range check.
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pg != 0
	}
	return In(r, GraphicRanges...)
}

// IsPrint reports whether the rune is defined as printable by Go. Such
// characters include letters, marks, numbers, punctuation, symbols, and the
// ASCII space character, from categories L, M, N, P, S and the ASCII space
// character.  This categorization is the same as IsGraphic except that the
// only spacing character is ASCII space, U+0020.
func IsPrint(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pp != 0
	}
	return In(r, PrintRanges...)
}

// IsOneOf reports whether the rune is a member of one of the ranges.
// The function "In" provides a nicer signature and should be used in preference to IsOneOf.
func IsOneOf(ranges []*RangeTable, r rune) bool {
	for _, inside := range ranges {
		if Is(inside, r) {
			return true
		}
	}
	return false
}

// In reports whether the rune is a member of one of the ranges.
func In(r rune, ranges ...*RangeTable) bool {
	for _, inside := range ranges {
		if Is(inside, r) {
			return true
		}
	}
	return false
}

// IsControl reports whether the rune is a control character.
// The C (Other) Unicode category includes more code points
// such as surrogates; use Is(C, r) to test for them.
func IsControl(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pC != 0
	}
	// All control characters are < MaxLatin1.
	return false
}

// IsLetter reports whether the rune is a letter (category L).
func IsLetter(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&(pLmask) != 0
	}
	return isExcludingLatin(Letter, r)
}

// IsMark reports whether the rune is a mark character (category M).
func IsMark(r rune) bool {
	// There are no mark characters in Latin-1.
	return isExcludingLatin(Mark, r)
}

// IsNumber reports whether the rune is a number (category N).
func IsNumber(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pN != 0
	}
	return isExcludingLatin(Number, r)
}

// IsPunct reports whether the rune is a Unicode punctuation character
// (category P).
func IsPunct(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pP != 0
	}
	return Is(Punct, r)
}

// IsSpace reports whether the rune is a space character as defined
// by Unicode's White Space property; in the Latin-1 space
// this is
//	'\t', '\n', '\v', '\f', '\r', ' ', U+0085 (NEL), U+00A0 (NBSP).
// Other definitions of spacing characters are set by category
// Z and property Pattern_White_Space.
func IsSpace(r rune) bool {
	// This property isn't the same as Z; special-case it.
	if uint32(r) <= MaxLatin1 {
		switch r {
		case '\t', '\n', '\v', '\f', '\r', ' ', 0x85, 0xA0:
			return true
		}
		return false
	}
	return isExcludingLatin(White_Space, r)
}

// IsSymbol reports whether the rune is a symbolic character.
func IsSymbol(r rune) bool {
	if uint32(r) <= MaxLatin1 {
		return properties[uint8(r)]&pS != 0
	}
	return isExcludingLatin(Symbol, r)
}
<|endoftext|>


<commit_before>package resolver

type MemStorage struct {
	Table    map[string]string
}

func NewMemStorage() *MemStorage {
    return &MemStorage{Table: make(map[string]string)}
}

func (this *MemStorage) Get(key string) string {
	return this.Table[key]
}

func (this *MemStorage) Set(key string, value string) {
	this.Table[key] = value
}

func (this *MemStorage) List() []string {
	var keys []string

	for k := range this.Table {
		keys = append(keys, k)
	}

	return keys
}
<commit_msg>oops. forgot formatting<commit_after>package resolver

type MemStorage struct {
	Table map[string]string
}

func NewMemStorage() *MemStorage {
	return &MemStorage{Table: make(map[string]string)}
}

func (this *MemStorage) Get(key string) string {
	return this.Table[key]
}

func (this *MemStorage) Set(key string, value string) {
	this.Table[key] = value
}

func (this *MemStorage) List() []string {
	var keys []string

	for k := range this.Table {
		keys = append(keys, k)
	}

	return keys
}
<|endoftext|>


<commit_before>// Copyright 2015 The Prometheus 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 discovery

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"time"

	"github.com/prometheus/common/log"
	"github.com/prometheus/common/model"
	"golang.org/x/net/context"
	"gopkg.in/fsnotify.v1"
	"gopkg.in/yaml.v2"

	"github.com/prometheus/prometheus/config"
)

const fileSDFilepathLabel = model.MetaLabelPrefix + "filepath"

// FileDiscovery provides service discovery functionality based
// on files that contain target groups in JSON or YAML format. Refreshing
// happens using file watches and periodic refreshes.
type FileDiscovery struct {
	paths    []string
	watcher  *fsnotify.Watcher
	interval time.Duration

	// lastRefresh stores which files were found during the last refresh
	// and how many target groups they contained.
	// This is used to detect deleted target groups.
	lastRefresh map[string]int
}

// NewFileDiscovery returns a new file discovery for the given paths.
func NewFileDiscovery(conf *config.FileSDConfig) *FileDiscovery {
	return &FileDiscovery{
		paths:    conf.Files,
		interval: time.Duration(conf.RefreshInterval),
	}
}

// listFiles returns a list of all files that match the configured patterns.
func (fd *FileDiscovery) listFiles() []string {
	var paths []string
	for _, p := range fd.paths {
		files, err := filepath.Glob(p)
		if err != nil {
			log.Errorf("Error expanding glob %q: %s", p, err)
			continue
		}
		paths = append(paths, files...)
	}
	return paths
}

// watchFiles sets watches on all full paths or directories that were configured for
// this file discovery.
func (fd *FileDiscovery) watchFiles() {
	if fd.watcher == nil {
		panic("no watcher configured")
	}
	for _, p := range fd.paths {
		if idx := strings.LastIndex(p, "/"); idx > -1 {
			p = p[:idx]
		} else {
			p = "./"
		}
		if err := fd.watcher.Add(p); err != nil {
			log.Errorf("Error adding file watch for %q: %s", p, err)
		}
	}
}

// Run implements the TargetProvider interface.
func (fd *FileDiscovery) Run(ctx context.Context, ch chan<- []*config.TargetGroup) {
	defer close(ch)
	defer fd.stop()

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Errorf("Error creating file watcher: %s", err)
		return
	}
	fd.watcher = watcher

	fd.refresh(ch)

	ticker := time.NewTicker(fd.interval)
	defer ticker.Stop()

	for {
		// Stopping has priority over refreshing. Thus we wrap the actual select
		// clause to always catch done signals.
		select {
		case <-ctx.Done():
			return
		default:
			select {
			case <-ctx.Done():
				return

			case event := <-fd.watcher.Events:
				// fsnotify sometimes sends a bunch of events without name or operation.
				// It's unclear what they are and why they are sent - filter them out.
				if len(event.Name) == 0 {
					break
				}
				// Everything but a chmod requires rereading.
				if event.Op^fsnotify.Chmod == 0 {
					break
				}
				// Changes to a file can spawn various sequences of events with
				// different combinations of operations. For all practical purposes
				// this is inaccurate.
				// The most reliable solution is to reload everything if anything happens.
				fd.refresh(ch)

			case <-ticker.C:
				// Setting a new watch after an update might fail. Make sure we don't lose
				// those files forever.
				fd.refresh(ch)

			case err := <-fd.watcher.Errors:
				if err != nil {
					log.Errorf("Error on file watch: %s", err)
				}
			}
		}
	}
}

// stop shuts down the file watcher.
func (fd *FileDiscovery) stop() {
	log.Debugf("Stopping file discovery for %s...", fd.paths)

	done := make(chan struct{})
	defer close(done)

	// Closing the watcher will deadlock unless all events and errors are drained.
	go func() {
		for {
			select {
			case <-fd.watcher.Errors:
			case <-fd.watcher.Events:
				// Drain all events and errors.
			case <-done:
				return
			}
		}
	}()
	if err := fd.watcher.Close(); err != nil {
		log.Errorf("Error closing file watcher for %s: %s", fd.paths, err)
	}

	log.Debugf("File discovery for %s stopped.", fd.paths)
}

// refresh reads all files matching the discovery's patterns and sends the respective
// updated target groups through the channel.
func (fd *FileDiscovery) refresh(ch chan<- []*config.TargetGroup) {
	ref := map[string]int{}
	for _, p := range fd.listFiles() {
		tgroups, err := readFile(p)
		if err != nil {
			log.Errorf("Error reading file %q: %s", p, err)
			// Prevent deletion down below.
			ref[p] = fd.lastRefresh[p]
			continue
		}
		ch <- tgroups

		ref[p] = len(tgroups)
	}
	// Send empty updates for sources that disappeared.
	for f, n := range fd.lastRefresh {
		m, ok := ref[f]
		if !ok || n > m {
			for i := m; i < n; i++ {
				ch <- []*config.TargetGroup{
					{Source: fileSource(f, i)},
				}
			}
		}
	}
	fd.lastRefresh = ref

	fd.watchFiles()
}

// fileSource returns a source ID for the i-th target group in the file.
func fileSource(filename string, i int) string {
	return fmt.Sprintf("%s:%d", filename, i)
}

// readFile reads a JSON or YAML list of targets groups from the file, depending on its
// file extension. It returns full configuration target groups.
func readFile(filename string) ([]*config.TargetGroup, error) {
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	var targetGroups []*config.TargetGroup

	switch ext := filepath.Ext(filename); strings.ToLower(ext) {
	case ".json":
		if err := json.Unmarshal(content, &targetGroups); err != nil {
			return nil, err
		}
	case ".yml", ".yaml":
		if err := yaml.Unmarshal(content, &targetGroups); err != nil {
			return nil, err
		}
	default:
		panic(fmt.Errorf("retrieval.FileDiscovery.readFile: unhandled file extension %q", ext))
	}

	for i, tg := range targetGroups {
		tg.Source = fileSource(filename, i)
		if tg.Labels == nil {
			tg.Labels = model.LabelSet{}
		}
		tg.Labels[fileSDFilepathLabel] = model.LabelValue(filename)
	}
	return targetGroups, nil
}
<commit_msg>Add File-SD metrics (#2103)<commit_after>// Copyright 2015 The Prometheus 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 discovery

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"github.com/prometheus/common/model"
	"golang.org/x/net/context"
	"gopkg.in/fsnotify.v1"
	"gopkg.in/yaml.v2"

	"github.com/prometheus/prometheus/config"
)

const fileSDFilepathLabel = model.MetaLabelPrefix + "filepath"

var (
	fileSDScanDuration = prometheus.NewSummary(
		prometheus.SummaryOpts{
			Namespace: namespace,
			Name:      "sd_file_scan_duration_seconds",
			Help:      "The duration of the File-SD scan in seconds.",
		})
	fileSDReadErrorsCount = prometheus.NewCounter(
		prometheus.CounterOpts{
			Namespace: namespace,
			Name:      "sd_file_read_errors_total",
			Help:      "The number of File-SD read errors.",
		})
)

func init() {
	prometheus.MustRegister(fileSDScanDuration)
	prometheus.MustRegister(fileSDReadErrorsCount)
}

// FileDiscovery provides service discovery functionality based
// on files that contain target groups in JSON or YAML format. Refreshing
// happens using file watches and periodic refreshes.
type FileDiscovery struct {
	paths    []string
	watcher  *fsnotify.Watcher
	interval time.Duration

	// lastRefresh stores which files were found during the last refresh
	// and how many target groups they contained.
	// This is used to detect deleted target groups.
	lastRefresh map[string]int
}

// NewFileDiscovery returns a new file discovery for the given paths.
func NewFileDiscovery(conf *config.FileSDConfig) *FileDiscovery {
	return &FileDiscovery{
		paths:    conf.Files,
		interval: time.Duration(conf.RefreshInterval),
	}
}

// listFiles returns a list of all files that match the configured patterns.
func (fd *FileDiscovery) listFiles() []string {
	var paths []string
	for _, p := range fd.paths {
		files, err := filepath.Glob(p)
		if err != nil {
			log.Errorf("Error expanding glob %q: %s", p, err)
			continue
		}
		paths = append(paths, files...)
	}
	return paths
}

// watchFiles sets watches on all full paths or directories that were configured for
// this file discovery.
func (fd *FileDiscovery) watchFiles() {
	if fd.watcher == nil {
		panic("no watcher configured")
	}
	for _, p := range fd.paths {
		if idx := strings.LastIndex(p, "/"); idx > -1 {
			p = p[:idx]
		} else {
			p = "./"
		}
		if err := fd.watcher.Add(p); err != nil {
			log.Errorf("Error adding file watch for %q: %s", p, err)
		}
	}
}

// Run implements the TargetProvider interface.
func (fd *FileDiscovery) Run(ctx context.Context, ch chan<- []*config.TargetGroup) {
	defer close(ch)
	defer fd.stop()

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Errorf("Error creating file watcher: %s", err)
		return
	}
	fd.watcher = watcher

	fd.refresh(ch)

	ticker := time.NewTicker(fd.interval)
	defer ticker.Stop()

	for {
		// Stopping has priority over refreshing. Thus we wrap the actual select
		// clause to always catch done signals.
		select {
		case <-ctx.Done():
			return
		default:
			select {
			case <-ctx.Done():
				return

			case event := <-fd.watcher.Events:
				// fsnotify sometimes sends a bunch of events without name or operation.
				// It's unclear what they are and why they are sent - filter them out.
				if len(event.Name) == 0 {
					break
				}
				// Everything but a chmod requires rereading.
				if event.Op^fsnotify.Chmod == 0 {
					break
				}
				// Changes to a file can spawn various sequences of events with
				// different combinations of operations. For all practical purposes
				// this is inaccurate.
				// The most reliable solution is to reload everything if anything happens.
				fd.refresh(ch)

			case <-ticker.C:
				// Setting a new watch after an update might fail. Make sure we don't lose
				// those files forever.
				fd.refresh(ch)

			case err := <-fd.watcher.Errors:
				if err != nil {
					log.Errorf("Error on file watch: %s", err)
				}
			}
		}
	}
}

// stop shuts down the file watcher.
func (fd *FileDiscovery) stop() {
	log.Debugf("Stopping file discovery for %s...", fd.paths)

	done := make(chan struct{})
	defer close(done)

	// Closing the watcher will deadlock unless all events and errors are drained.
	go func() {
		for {
			select {
			case <-fd.watcher.Errors:
			case <-fd.watcher.Events:
				// Drain all events and errors.
			case <-done:
				return
			}
		}
	}()
	if err := fd.watcher.Close(); err != nil {
		log.Errorf("Error closing file watcher for %s: %s", fd.paths, err)
	}

	log.Debugf("File discovery for %s stopped.", fd.paths)
}

// refresh reads all files matching the discovery's patterns and sends the respective
// updated target groups through the channel.
func (fd *FileDiscovery) refresh(ch chan<- []*config.TargetGroup) {
	t0 := time.Now()
	defer func() {
		fileSDScanDuration.Observe(time.Since(t0).Seconds())
	}()

	ref := map[string]int{}
	for _, p := range fd.listFiles() {
		tgroups, err := readFile(p)
		if err != nil {
			fileSDReadErrorsCount.Inc()
			log.Errorf("Error reading file %q: %s", p, err)
			// Prevent deletion down below.
			ref[p] = fd.lastRefresh[p]
			continue
		}
		ch <- tgroups

		ref[p] = len(tgroups)
	}
	// Send empty updates for sources that disappeared.
	for f, n := range fd.lastRefresh {
		m, ok := ref[f]
		if !ok || n > m {
			for i := m; i < n; i++ {
				ch <- []*config.TargetGroup{
					{Source: fileSource(f, i)},
				}
			}
		}
	}
	fd.lastRefresh = ref

	fd.watchFiles()
}

// fileSource returns a source ID for the i-th target group in the file.
func fileSource(filename string, i int) string {
	return fmt.Sprintf("%s:%d", filename, i)
}

// readFile reads a JSON or YAML list of targets groups from the file, depending on its
// file extension. It returns full configuration target groups.
func readFile(filename string) ([]*config.TargetGroup, error) {
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, err
	}

	var targetGroups []*config.TargetGroup

	switch ext := filepath.Ext(filename); strings.ToLower(ext) {
	case ".json":
		if err := json.Unmarshal(content, &targetGroups); err != nil {
			return nil, err
		}
	case ".yml", ".yaml":
		if err := yaml.Unmarshal(content, &targetGroups); err != nil {
			return nil, err
		}
	default:
		panic(fmt.Errorf("retrieval.FileDiscovery.readFile: unhandled file extension %q", ext))
	}

	for i, tg := range targetGroups {
		tg.Source = fileSource(filename, i)
		if tg.Labels == nil {
			tg.Labels = model.LabelSet{}
		}
		tg.Labels[fileSDFilepathLabel] = model.LabelValue(filename)
	}
	return targetGroups, nil
}
<|endoftext|>


<commit_before>package lucene42

import (
	"github.com/balzaczyy/golucene/core/store"
	"testing"
)

func TestReadFieldInfos(t *testing.T) {
	path := "../../search/testdata/osx/belfrysample"
	d, err := store.OpenFSDirectory(path)
	if err != nil {
		t.Error(err)
	}
	ctx := store.NewIOContextBool(false)
	cd, err := store.NewCompoundFileDirectory(d, "_0.cfs", ctx, false)
	if err != nil {
		t.Error(err)
	}
	fis, err := Lucene42FieldInfosReader(cd, "_0", store.IO_CONTEXT_READONCE)
	if err != nil {
		t.Error(err)
	}
	if !fis.HasNorms || fis.HasDocValues {
		t.Errorf("hasNorms must be true and hasDocValues must be false, but found %v", fis)
	}
}
<commit_msg>fix unit test for codec/lucene42<commit_after>package lucene42

import (
	"github.com/balzaczyy/golucene/core/store"
	"testing"
)

func TestReadFieldInfos(t *testing.T) {
	path := "../../search/testdata/osx/belfrysample"
	d, err := store.OpenFSDirectory(path)
	if err != nil {
		t.Error(err)
	}
	ctx := store.NewIOContextBool(false)
	cd, err := store.NewCompoundFileDirectory(d, "_0.cfs", ctx, false)
	if err != nil {
		t.Error(err)
	}
	fis, err := Lucene42FieldInfosReader(cd, "_0", "", store.IO_CONTEXT_READONCE)
	if err != nil {
		t.Error(err)
	}
	if !fis.HasNorms || fis.HasDocValues {
		t.Errorf("hasNorms must be true and hasDocValues must be false, but found %v", fis)
	}
}
<|endoftext|>