<commit_before>/*
 *  Charon: A game authentication server
 *  Copyright (C) 2014-2015  Alex Mayfield <alexmax2742@gmail.com>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package charon

import (
	"crypto/sha256"
	"strings"
	"sync"

	"github.com/jmoiron/sqlx"
	_ "github.com/mattn/go-sqlite3"
	"github.com/tadglines/go-pkgs/crypto/srp"
)

type Database struct {
	db    *sqlx.DB
	mutex sync.Mutex
}

// Schema for sqlite3.
var schema = `
CREATE TABLE IF NOT EXISTS users(
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	username TEXT UNIQUE NOT NULL,
	email TEXT NOT NULL,
	verifier BLOB NOT NULL,
	salt BLOB NOT NULL,
	access TEXT NOT NULL,
	active INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS profiles(
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	user_id INTEGER,
	clan TEXT,
	contactinfo TEXT,
	country TEXT,
	gravatar TEXT,
	location TEXT,
	message TEXT,
	username TEXT,
	visible INTEGER,
	visible_lastplayed INTEGER,
	FOREIGN KEY(user_id) REFERENCES users(id)
);`

func NewDatabase() (database *Database, err error) {
	// Create a database connection.
	db, err := sqlx.Connect("sqlite3", ":memory:")
	if err != nil {
		return
	}

	// Create the database schema.
	_ = db.MustExec("PRAGMA foreign_keys = ON;")
	_ = db.MustExec(schema)

	database = new(Database)
	database.db = db
	return
}

// A representation of a user.
type User struct {
	Id       uint
	Username string
	Email    string
	Verifier []byte
	Salt     []byte
	Access   string
	Active   bool
}

const (
	UserAccessUnverified string = "UNVERIFIED"
	UserAccessUser       string = "USER"
	UserAccessOp         string = "OP"
	UserAccessMaster     string = "MASTER"
	UserAccessOwner      string = "OWNER"
)

// Add a new user
func (self *Database) AddUser(username string, email string, password string) (err error) {
	srp, err := srp.NewSRP("rfc5054.2048", sha256.New, nil)
	if err != nil {
		return err
	}

	user := new(User)
	user.Username = username
	user.Email = email
	user.Access = UserAccessUnverified
	user.Active = false

	user.Salt, user.Verifier, err = srp.ComputeVerifier([]byte(password))
	if err != nil {
		return err
	}

	self.mutex.Lock()
	_, err = self.db.NamedExec("INSERT INTO users (Username, Email, Verifier, Salt, Access, Active) VALUES (:username, :email, :verifier, :salt, :access, :active)", user)
	self.mutex.Unlock()
	return
}

// Try to find a specific user by name or email address.
func (self *Database) FindUser(username string) (user *User, err error) {
	user = &User{}
	self.mutex.Lock()
	err = self.db.Get(user, "SELECT * FROM users WHERE username LIKE $1 OR email LIKE $1", strings.ToLower(username))
	self.mutex.Unlock()
	return
}

// A representation of a user's profile.
type Profile struct {
	Id                 uint
	User_id            uint
	Clan               string
	Contactinfo        string
	Country            string
	Gravatar           string
	Location           string
	Message            string
	Username           string
	Visible            bool
	Visible_lastplayed bool
}
<commit_msg>Linter issues.<commit_after>/*
 *  Charon: A game authentication server
 *  Copyright (C) 2014-2015  Alex Mayfield <alexmax2742@gmail.com>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package charon

import (
	"crypto/sha256"
	"strings"
	"sync"

	"github.com/jmoiron/sqlx"
	_ "github.com/mattn/go-sqlite3" // Database driver
	"github.com/tadglines/go-pkgs/crypto/srp"
)

// Database is an instance of our database connection and all necessary state
// used to manage said instance.
type Database struct {
	db    *sqlx.DB
	mutex sync.Mutex
}

// Schema for sqlite3.
const schema = `
CREATE TABLE IF NOT EXISTS users(
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	username TEXT UNIQUE NOT NULL,
	email TEXT NOT NULL,
	verifier BLOB NOT NULL,
	salt BLOB NOT NULL,
	access TEXT NOT NULL,
	active INTEGER NOT NULL
);
CREATE TABLE IF NOT EXISTS profiles(
	id INTEGER PRIMARY KEY AUTOINCREMENT,
	user_id INTEGER,
	clan TEXT,
	contactinfo TEXT,
	country TEXT,
	gravatar TEXT,
	location TEXT,
	message TEXT,
	username TEXT,
	visible INTEGER,
	visible_lastplayed INTEGER,
	FOREIGN KEY(user_id) REFERENCES users(id)
);`

// NewDatabase creates a new Database instance.
func NewDatabase() (database *Database, err error) {
	// Create a database connection.
	db, err := sqlx.Connect("sqlite3", ":memory:")
	if err != nil {
		return
	}

	// Create the database schema.
	_ = db.MustExec("PRAGMA foreign_keys = ON;")
	_ = db.MustExec(schema)

	database = new(Database)
	database.db = db
	return
}

// User is a representation of the `user` table in the database.
type User struct {
	ID       uint
	Username string
	Email    string
	Verifier []byte
	Salt     []byte
	Access   string
	Active   bool
}

// User access constants.
const (
	UserAccessUnverified string = "UNVERIFIED"
	UserAccessUser       string = "USER"
	UserAccessOp         string = "OP"
	UserAccessMaster     string = "MASTER"
	UserAccessOwner      string = "OWNER"
)

// AddUser adds a new user.
func (database *Database) AddUser(username string, email string, password string) (err error) {
	srp, err := srp.NewSRP("rfc5054.2048", sha256.New, nil)
	if err != nil {
		return err
	}

	user := new(User)
	user.Username = username
	user.Email = email
	user.Access = UserAccessUnverified
	user.Active = false

	user.Salt, user.Verifier, err = srp.ComputeVerifier([]byte(password))
	if err != nil {
		return err
	}

	database.mutex.Lock()
	_, err = database.db.NamedExec("INSERT INTO users (Username, Email, Verifier, Salt, Access, Active) VALUES (:username, :email, :verifier, :salt, :access, :active)", user)
	database.mutex.Unlock()
	return
}

// FindUser tries to find a specific user by name or email address.
func (database *Database) FindUser(username string) (user *User, err error) {
	user = &User{}
	database.mutex.Lock()
	err = database.db.Get(user, "SELECT * FROM users WHERE username LIKE $1 OR email LIKE $1", strings.ToLower(username))
	database.mutex.Unlock()
	return
}

// Profile is representation of the `profile` table in the database.
type Profile struct {
	ID                 uint
	User_id            uint
	Clan               string
	Contactinfo        string
	Country            string
	Gravatar           string
	Location           string
	Message            string
	Username           string
	Visible            bool
	Visible_lastplayed bool
}
<|endoftext|>


<commit_before>package gramework

import (
	"github.com/valyala/fasthttp"
)

// ListenAndServe on given addr
func (app *App) ListenAndServe(addr string) error {
	return fasthttp.ListenAndServe(addr, app.router.Handler)
}
<commit_msg>app:listenAndServe: make address argument optional<commit_after>package gramework

import (
	"errors"
	"flag"

	"github.com/valyala/fasthttp"
)

// ListenAndServe on given addr
func (app *App) ListenAndServe(addr ...string) error {
	var bind string
	if len(addr) > 0 {
		bind = addr[0]
	} else {
		if !app.flagsRegistered {
			app.RegFlags()
		}
	}
	if !flag.Parsed() {
		flag.Parse()
	}
	if app.Flags.values != nil {
		if bindFlag, ok := app.Flags.values["bind"]; ok {
			bind = *bindFlag.Value
		}
	}
	if bind == "" {
		return errors.New("No bind address provided")
	}
	l := app.Logger.WithField("bind", bind)
	l.Info("Starting")
	err := fasthttp.ListenAndServe(bind, app.router.Handler)
	l.Errorf("ListenAndServe failed: %s", err)
	return err
}
<|endoftext|>


<commit_before>package h2spec

import (
	"bytes"
	"github.com/bradfitz/http2"
	"github.com/bradfitz/http2/hpack"
	"time"
)

func TestData(ctx *Context) {
	PrintHeader("6.1. DATA", 0)

	func(ctx *Context) {
		desc := "Sends a DATA frame with 0x0 stream identifier"
		msg := "The endpoint MUST respond with a connection error of type PROTOCOL_ERROR."
		result := false

		http2Conn := CreateHttp2Conn(ctx, true)
		defer http2Conn.conn.Close()

		http2Conn.fr.WriteData(0, true, []byte("test"))
		timeCh := time.After(3 * time.Second)

	loop:
		for {
			select {
			case f := <-http2Conn.dataCh:
				gf, ok := f.(*http2.GoAwayFrame)
				if ok {
					if gf.ErrCode == http2.ErrCodeProtocol {
						result = true
					}
				}
			case <-http2Conn.errCh:
				break loop
			case <-timeCh:
				break loop
			}
		}

		PrintResult(result, desc, msg, 0)
	}(ctx)

	func(ctx *Context) {
		desc := "Sends a DATA frame on the stream that is not opend"
		msg := "The endpoint MUST respond with a stream error of type STREAM_CLOSED."
		result := false

		http2Conn := CreateHttp2Conn(ctx, true)
		defer http2Conn.conn.Close()

		var buf bytes.Buffer
		hdrs := []hpack.HeaderField{
			pair(":method", "GET"),
			pair(":scheme", "http"),
			pair(":path", "/"),
			pair(":authority", ctx.Authority()),
		}
		enc := hpack.NewEncoder(&buf)
		for _, hf := range hdrs {
			_ = enc.WriteField(hf)
		}

		var hp http2.HeadersFrameParam
		hp.StreamID = 1
		hp.EndStream = true
		hp.EndHeaders = true
		hp.BlockFragment = buf.Bytes()
		http2Conn.fr.WriteHeaders(hp)
		http2Conn.fr.WriteData(1, true, []byte("test"))

		timeCh := time.After(3 * time.Second)

	loop:
		for {
			select {
			case f := <-http2Conn.dataCh:
				gf, ok := f.(*http2.GoAwayFrame)
				if ok {
					if gf.ErrCode == http2.ErrCodeStreamClosed {
						result = true
					}
				}
			case <-http2Conn.errCh:
				break loop
			case <-timeCh:
				break loop
			}
		}

		PrintResult(result, desc, msg, 0)
	}(ctx)

	// bradfitz/http2 does not support padding of DATA frame.
	/*
		func(ctx *Context) {
			desc := "Sends a DATA frame with invalid pad length"
			msg := "The endpoint MUST treat this as a connection error of type PROTOCOL_ERROR."
			result := false

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			var buf bytes.Buffer
			hdrs := []hpack.HeaderField{
				pair(":method", "GET"),
				pair(":scheme", "http"),
				pair(":path", "/"),
				pair(":authority", ctx.Authority()),
			}
			enc := hpack.NewEncoder(&buf)
			for _, hf := range hdrs {
				_ = enc.WriteField(hf)
			}

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = true
			hp.EndHeaders = true
			hp.BlockFragment = buf.Bytes()
			http2Conn.fr.WriteHeaders(hp)

			http2Conn.fr.WriteData(1, true, []byte("test"))

			timeCh := time.After(3 * time.Second)

		loop:
			for {
				select {
				case f := <-http2Conn.dataCh:
					gf, ok := f.(*http2.GoAwayFrame)
					if ok {
						if gf.ErrCode == http2.ErrCodeProtocol {
							result = true
						}
					}
				case <-http2Conn.errCh:
					break loop
				case <-timeCh:
					break loop
				}
			}

			PrintResult(result, desc, msg, 0)
		}(ctx)
	*/

	PrintFooter()
}
<commit_msg>should check against RST_STREAM<commit_after>package h2spec

import (
	"bytes"
	"github.com/bradfitz/http2"
	"github.com/bradfitz/http2/hpack"
	"time"
)

func TestData(ctx *Context) {
	PrintHeader("6.1. DATA", 0)

	func(ctx *Context) {
		desc := "Sends a DATA frame with 0x0 stream identifier"
		msg := "The endpoint MUST respond with a connection error of type PROTOCOL_ERROR."
		result := false

		http2Conn := CreateHttp2Conn(ctx, true)
		defer http2Conn.conn.Close()

		http2Conn.fr.WriteData(0, true, []byte("test"))
		timeCh := time.After(3 * time.Second)

	loop:
		for {
			select {
			case f := <-http2Conn.dataCh:
				gf, ok := f.(*http2.GoAwayFrame)
				if ok {
					if gf.ErrCode == http2.ErrCodeProtocol {
						result = true
					}
				}
			case <-http2Conn.errCh:
				break loop
			case <-timeCh:
				break loop
			}
		}

		PrintResult(result, desc, msg, 0)
	}(ctx)

	func(ctx *Context) {
		desc := "Sends a DATA frame on the stream that is not opend"
		msg := "The endpoint MUST respond with a stream error of type STREAM_CLOSED."
		result := false

		http2Conn := CreateHttp2Conn(ctx, true)
		defer http2Conn.conn.Close()

		var buf bytes.Buffer
		hdrs := []hpack.HeaderField{
			pair(":method", "GET"),
			pair(":scheme", "http"),
			pair(":path", "/"),
			pair(":authority", ctx.Authority()),
		}
		enc := hpack.NewEncoder(&buf)
		for _, hf := range hdrs {
			_ = enc.WriteField(hf)
		}

		var hp http2.HeadersFrameParam
		hp.StreamID = 1
		hp.EndStream = true
		hp.EndHeaders = true
		hp.BlockFragment = buf.Bytes()
		http2Conn.fr.WriteHeaders(hp)
		http2Conn.fr.WriteData(1, true, []byte("test"))

		timeCh := time.After(3 * time.Second)

	loop:
		for {
			select {
			case f := <-http2Conn.dataCh:
				gf, ok := f.(*http2.RSTStreamFrame)
				if ok {
					if gf.ErrCode == http2.ErrCodeStreamClosed {
						result = true
					}
				}
			case <-http2Conn.errCh:
				break loop
			case <-timeCh:
				break loop
			}
		}

		PrintResult(result, desc, msg, 0)
	}(ctx)

	// bradfitz/http2 does not support padding of DATA frame.
	/*
		func(ctx *Context) {
			desc := "Sends a DATA frame with invalid pad length"
			msg := "The endpoint MUST treat this as a connection error of type PROTOCOL_ERROR."
			result := false

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			var buf bytes.Buffer
			hdrs := []hpack.HeaderField{
				pair(":method", "GET"),
				pair(":scheme", "http"),
				pair(":path", "/"),
				pair(":authority", ctx.Authority()),
			}
			enc := hpack.NewEncoder(&buf)
			for _, hf := range hdrs {
				_ = enc.WriteField(hf)
			}

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = true
			hp.EndHeaders = true
			hp.BlockFragment = buf.Bytes()
			http2Conn.fr.WriteHeaders(hp)

			http2Conn.fr.WriteData(1, true, []byte("test"))

			timeCh := time.After(3 * time.Second)

		loop:
			for {
				select {
				case f := <-http2Conn.dataCh:
					gf, ok := f.(*http2.GoAwayFrame)
					if ok {
						if gf.ErrCode == http2.ErrCodeProtocol {
							result = true
						}
					}
				case <-http2Conn.errCh:
					break loop
				case <-timeCh:
					break loop
				}
			}

			PrintResult(result, desc, msg, 0)
		}(ctx)
	*/

	PrintFooter()
}
<|endoftext|>


<commit_before>package h2spec

import (
	"errors"
	"fmt"
	"golang.org/x/net/http2"
	"io"
	"net"
	"syscall"
)

func WindowUpdateTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9", "WINDOW_UPDATE")

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame",
		"The endpoint is expected to send the DATA frame based on the window size.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			pass = false
			expected = []Result{
				&ResultFrame{LengthDefault, http2.FrameData, FlagDefault, ErrCodeDefault},
			}

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			settings := http2.Setting{http2.SettingInitialWindowSize, 1}
			http2Conn.fr.WriteSettings(settings)

			hdrs := commonHeaderFields(ctx)

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = true
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)

			winUpdated := false

		loop:
			for {
				f, err := http2Conn.ReadFrame(ctx.Timeout)
				if err != nil {
					opErr, ok := err.(*net.OpError)
					if err == io.EOF || (ok && opErr.Err == syscall.ECONNRESET) {
						rf, ok := actual.(*ResultFrame)
						if actual == nil || (ok && rf.Type != http2.FrameGoAway) {
							actual = &ResultConnectionClose{}
						}
					} else if err == TIMEOUT {
						if actual == nil {
							actual = &ResultTestTimeout{}
						}
					} else {
						actual = &ResultError{err}
					}
					break loop
				}

				switch f := f.(type) {
				case *http2.DataFrame:
					if winUpdated {
						if f.FrameHeader.Length > 10 {
							err := errors.New("The length of DATA frame is invalid.")
							actual = &ResultError{err}
							break loop
						}

						actual = CreateResultFrame(f)
						pass = true
						break loop
					} else {
						if f.FrameHeader.Length != 1 {
							err := errors.New("The length of DATA frame is invalid.")
							actual = &ResultError{err}
							break loop
						}

						http2Conn.fr.WriteWindowUpdate(1, 10)
						winUpdated = true
					}
				case *http2.GoAwayFrame:
					actual = CreateResultFrame(f)
					break loop
				case *http2.RSTStreamFrame:
					actual = CreateResultFrame(f)
					break loop
				default:
					actual = CreateResultFrame(f)
				}
			}

			return pass, expected, actual
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with an flow control window increment of 0",
		"The endpoint MUST respond with a connection error of type PROTOCOL_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			http2Conn.fr.WriteWindowUpdate(0, 0)

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with an flow control window increment of 0 on a stream",
		"The endpoint MUST respond with a stream error of type PROTOCOL_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = false
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)
			http2Conn.fr.WriteWindowUpdate(1, 0)

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestStreamError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with a length other than a multiple of 4 octets",
		"The endpoint MUST respond with a connection error of type FRAME_SIZE_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			fmt.Fprintf(http2Conn.conn, "\x00\x00\x03\x08\x00\x00\x00\x00\x00")
			fmt.Fprintf(http2Conn.conn, "\x00\x00\x01")

			actualCodes := []http2.ErrCode{http2.ErrCodeFrameSize}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestGroup(TheFlowControlWindowTestGroup(ctx))
	tg.AddTestGroup(InitialFlowControlWindowSizeTestGroup(ctx))

	return tg
}

func TheFlowControlWindowTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9.1", "The Flow Control Window")

	tg.AddTestCase(NewTestCase(
		"Sends multiple WINDOW_UPDATE frames on a connection increasing the flow control window to above 2^31-1",
		"The endpoint MUST sends a GOAWAY frame with a FLOW_CONTROL_ERROR code.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			pass = false
			expected = []Result{
				&ResultFrame{LengthDefault, http2.FrameGoAway, FlagDefault, http2.ErrCodeFlowControl},
			}

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			http2Conn.fr.WriteWindowUpdate(0, 2147483647)
			http2Conn.fr.WriteWindowUpdate(0, 2147483647)

		loop:
			for {
				f, err := http2Conn.ReadFrame(ctx.Timeout)
				if err != nil {
					opErr, ok := err.(*net.OpError)
					if err == io.EOF || (ok && opErr.Err == syscall.ECONNRESET) {
						rf, ok := actual.(*ResultFrame)
						if actual == nil || (ok && rf.Type != http2.FrameGoAway) {
							actual = &ResultConnectionClose{}
						}
					} else if err == TIMEOUT {
						if actual == nil {
							actual = &ResultTestTimeout{}
						}
					} else {
						actual = &ResultError{err}
					}
					break loop
				}

				switch f := f.(type) {
				case *http2.GoAwayFrame:
					actual = CreateResultFrame(f)
					if f.ErrCode == http2.ErrCodeFlowControl {
						pass = true
					}
					break loop
				default:
					actual = CreateResultFrame(f)
				}
			}

			return pass, expected, actual
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends multiple WINDOW_UPDATE frames on a stream increasing the flow control window to above 2^31-1",
		"The endpoint MUST sends a RST_STREAM with the error code of FLOW_CONTROL_ERROR code.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			pass = false
			expected = []Result{
				&ResultFrame{LengthDefault, http2.FrameRSTStream, FlagDefault, http2.ErrCodeFlowControl},
			}

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = false
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)

			http2Conn.fr.WriteWindowUpdate(1, 2147483647)
			http2Conn.fr.WriteWindowUpdate(1, 2147483647)

		loop:
			for {
				f, err := http2Conn.ReadFrame(ctx.Timeout)
				if err != nil {
					opErr, ok := err.(*net.OpError)
					if err == io.EOF || (ok && opErr.Err == syscall.ECONNRESET) {
						rf, ok := actual.(*ResultFrame)
						if actual == nil || (ok && rf.Type != http2.FrameGoAway) {
							actual = &ResultConnectionClose{}
						}
					} else if err == TIMEOUT {
						if actual == nil {
							actual = &ResultTestTimeout{}
						}
					} else {
						actual = &ResultError{err}
					}
					break loop
				}

				switch f := f.(type) {
				case *http2.GoAwayFrame:
					actual = CreateResultFrame(f)
					break loop
				case *http2.RSTStreamFrame:
					actual = CreateResultFrame(f)
					if f.ErrCode == http2.ErrCodeFlowControl {
						pass = true
					}
					break loop
				default:
					actual = CreateResultFrame(f)
				}
			}

			return pass, expected, actual
		},
	))

	return tg
}

func InitialFlowControlWindowSizeTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9.2", "Initial Flow Control Window Size")

	tg.AddTestCase(NewTestCase(
		"Sends a SETTINGS_INITIAL_WINDOW_SIZE settings with an exceeded maximum window size value",
		"The endpoint MUST respond with a connection error of type FLOW_CONTROL_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			fmt.Fprintf(http2Conn.conn, "\x00\x00\x06\x04\x00\x00\x00\x00\x00")
			fmt.Fprintf(http2Conn.conn, "\x00\x04\x80\x00\x00\x00")

			actualCodes := []http2.ErrCode{http2.ErrCodeFlowControl}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	return tg
}
<commit_msg>Skip window update test if DATA frame has END_STREAM flag<commit_after>package h2spec

import (
	"errors"
	"fmt"
	"golang.org/x/net/http2"
	"io"
	"net"
	"syscall"
)

func WindowUpdateTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9", "WINDOW_UPDATE")

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame",
		"The endpoint is expected to send the DATA frame based on the window size.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			pass = false
			expected = []Result{
				&ResultFrame{LengthDefault, http2.FrameData, FlagDefault, ErrCodeDefault},
			}

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			settings := http2.Setting{http2.SettingInitialWindowSize, 1}
			http2Conn.fr.WriteSettings(settings)

			hdrs := commonHeaderFields(ctx)

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = true
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)

			winUpdated := false

		loop:
			for {
				f, err := http2Conn.ReadFrame(ctx.Timeout)
				if err != nil {
					opErr, ok := err.(*net.OpError)
					if err == io.EOF || (ok && opErr.Err == syscall.ECONNRESET) {
						rf, ok := actual.(*ResultFrame)
						if actual == nil || (ok && rf.Type != http2.FrameGoAway) {
							actual = &ResultConnectionClose{}
						}
					} else if err == TIMEOUT {
						if actual == nil {
							actual = &ResultTestTimeout{}
						}
					} else {
						actual = &ResultError{err}
					}
					break loop
				}

				switch f := f.(type) {
				case *http2.DataFrame:
					if winUpdated {
						// Let's skip this test if the DATA frame has END_STREAM flag.
						if f.FrameHeader.Flags.Has(http2.FlagDataEndStream) {
							actual = &ResultSkipped{"The length of DATA frame is 0."}
							return true, nil, actual
						}

						if f.FrameHeader.Length != 1 {
							err := errors.New("The length of DATA frame is invalid.")
							actual = &ResultError{err}
							break loop
						}

						actual = CreateResultFrame(f)
						pass = true
						break loop
					} else {
						// Let's skip this test if the DATA frame has END_STREAM flag.
						if f.FrameHeader.Flags.Has(http2.FlagDataEndStream) {
							actual = &ResultSkipped{"The length of DATA frame is 0."}
							return true, nil, actual
						}

						if f.FrameHeader.Length != 1 {
							err := errors.New("The length of DATA frame is invalid.")
							actual = &ResultError{err}
							break loop
						}

						http2Conn.fr.WriteWindowUpdate(1, 1)
						winUpdated = true
					}
				case *http2.GoAwayFrame:
					actual = CreateResultFrame(f)
					break loop
				case *http2.RSTStreamFrame:
					actual = CreateResultFrame(f)
					break loop
				default:
					actual = CreateResultFrame(f)
				}
			}

			return pass, expected, actual
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with an flow control window increment of 0",
		"The endpoint MUST respond with a connection error of type PROTOCOL_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			http2Conn.fr.WriteWindowUpdate(0, 0)

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with an flow control window increment of 0 on a stream",
		"The endpoint MUST respond with a stream error of type PROTOCOL_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = false
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)
			http2Conn.fr.WriteWindowUpdate(1, 0)

			actualCodes := []http2.ErrCode{http2.ErrCodeProtocol}
			return TestStreamError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends a WINDOW_UPDATE frame with a length other than a multiple of 4 octets",
		"The endpoint MUST respond with a connection error of type FRAME_SIZE_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			fmt.Fprintf(http2Conn.conn, "\x00\x00\x03\x08\x00\x00\x00\x00\x00")
			fmt.Fprintf(http2Conn.conn, "\x00\x00\x01")

			actualCodes := []http2.ErrCode{http2.ErrCodeFrameSize}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	tg.AddTestGroup(TheFlowControlWindowTestGroup(ctx))
	tg.AddTestGroup(InitialFlowControlWindowSizeTestGroup(ctx))

	return tg
}

func TheFlowControlWindowTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9.1", "The Flow Control Window")

	tg.AddTestCase(NewTestCase(
		"Sends multiple WINDOW_UPDATE frames on a connection increasing the flow control window to above 2^31-1",
		"The endpoint MUST sends a GOAWAY frame with a FLOW_CONTROL_ERROR code.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			pass = false
			expected = []Result{
				&ResultFrame{LengthDefault, http2.FrameGoAway, FlagDefault, http2.ErrCodeFlowControl},
			}

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			http2Conn.fr.WriteWindowUpdate(0, 2147483647)
			http2Conn.fr.WriteWindowUpdate(0, 2147483647)

		loop:
			for {
				f, err := http2Conn.ReadFrame(ctx.Timeout)
				if err != nil {
					opErr, ok := err.(*net.OpError)
					if err == io.EOF || (ok && opErr.Err == syscall.ECONNRESET) {
						rf, ok := actual.(*ResultFrame)
						if actual == nil || (ok && rf.Type != http2.FrameGoAway) {
							actual = &ResultConnectionClose{}
						}
					} else if err == TIMEOUT {
						if actual == nil {
							actual = &ResultTestTimeout{}
						}
					} else {
						actual = &ResultError{err}
					}
					break loop
				}

				switch f := f.(type) {
				case *http2.GoAwayFrame:
					actual = CreateResultFrame(f)
					if f.ErrCode == http2.ErrCodeFlowControl {
						pass = true
					}
					break loop
				default:
					actual = CreateResultFrame(f)
				}
			}

			return pass, expected, actual
		},
	))

	tg.AddTestCase(NewTestCase(
		"Sends multiple WINDOW_UPDATE frames on a stream increasing the flow control window to above 2^31-1",
		"The endpoint MUST sends a RST_STREAM with the error code of FLOW_CONTROL_ERROR code.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			pass = false
			expected = []Result{
				&ResultFrame{LengthDefault, http2.FrameRSTStream, FlagDefault, http2.ErrCodeFlowControl},
			}

			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			hdrs := commonHeaderFields(ctx)

			var hp http2.HeadersFrameParam
			hp.StreamID = 1
			hp.EndStream = false
			hp.EndHeaders = true
			hp.BlockFragment = http2Conn.EncodeHeader(hdrs)
			http2Conn.fr.WriteHeaders(hp)

			http2Conn.fr.WriteWindowUpdate(1, 2147483647)
			http2Conn.fr.WriteWindowUpdate(1, 2147483647)

		loop:
			for {
				f, err := http2Conn.ReadFrame(ctx.Timeout)
				if err != nil {
					opErr, ok := err.(*net.OpError)
					if err == io.EOF || (ok && opErr.Err == syscall.ECONNRESET) {
						rf, ok := actual.(*ResultFrame)
						if actual == nil || (ok && rf.Type != http2.FrameGoAway) {
							actual = &ResultConnectionClose{}
						}
					} else if err == TIMEOUT {
						if actual == nil {
							actual = &ResultTestTimeout{}
						}
					} else {
						actual = &ResultError{err}
					}
					break loop
				}

				switch f := f.(type) {
				case *http2.GoAwayFrame:
					actual = CreateResultFrame(f)
					break loop
				case *http2.RSTStreamFrame:
					actual = CreateResultFrame(f)
					if f.ErrCode == http2.ErrCodeFlowControl {
						pass = true
					}
					break loop
				default:
					actual = CreateResultFrame(f)
				}
			}

			return pass, expected, actual
		},
	))

	return tg
}

func InitialFlowControlWindowSizeTestGroup(ctx *Context) *TestGroup {
	tg := NewTestGroup("6.9.2", "Initial Flow Control Window Size")

	tg.AddTestCase(NewTestCase(
		"Sends a SETTINGS_INITIAL_WINDOW_SIZE settings with an exceeded maximum window size value",
		"The endpoint MUST respond with a connection error of type FLOW_CONTROL_ERROR.",
		func(ctx *Context) (pass bool, expected []Result, actual Result) {
			http2Conn := CreateHttp2Conn(ctx, true)
			defer http2Conn.conn.Close()

			fmt.Fprintf(http2Conn.conn, "\x00\x00\x06\x04\x00\x00\x00\x00\x00")
			fmt.Fprintf(http2Conn.conn, "\x00\x04\x80\x00\x00\x00")

			actualCodes := []http2.ErrCode{http2.ErrCodeFlowControl}
			return TestConnectionError(ctx, http2Conn, actualCodes)
		},
	))

	return tg
}
<|endoftext|>


<commit_before>// Copyright 2014 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 reader // import "golang.org/x/tour/reader"

import (
	"fmt"
	"io"
	"os"
)

func Validate(r io.Reader) {
	b := make([]byte, 1024)
	i, o := 0, 0
	for ; i < 1<<20 && o < 1<<20; i++ { // test 1mb
		n, err := r.Read(b)
		for i, v := range b[:n] {
			if v != 'A' {
				fmt.Fprintf(os.Stderr, "got byte %x at offset %v, want 'A'\n", v, o+i)
				return
			}
		}
		o += n
		if err != nil {
			fmt.Fprintf(os.Stderr, "read error: %v\n", err)
			return
		}
	}
	if o == 0 {
		fmt.Fprintf(os.Stderr, "read zero bytes after %d Read calls\n", i)
		return
	}
	fmt.Println("OK!")
}
<commit_msg>reader: move error check in validate<commit_after>// Copyright 2014 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 reader // import "golang.org/x/tour/reader"

import (
	"fmt"
	"io"
	"os"
)

func Validate(r io.Reader) {
	b := make([]byte, 1024)
	i, o := 0, 0
	for ; i < 1<<20 && o < 1<<20; i++ { // test 1mb
		n, err := r.Read(b)
		if err != nil {
			fmt.Fprintf(os.Stderr, "read error: %v\n", err)
			return
		}
		for i, v := range b[:n] {
			if v != 'A' {
				fmt.Fprintf(os.Stderr, "got byte %x at offset %v, want 'A'\n", v, o+i)
				return
			}
		}
		o += n
	}
	if o == 0 {
		fmt.Fprintf(os.Stderr, "read zero bytes after %d Read calls\n", i)
		return
	}
	fmt.Println("OK!")
}
<|endoftext|>


<commit_before>// Licensed under the Apache License, Version 2.0
// Details: https://raw.githubusercontent.com/maniksurtani/quotaservice/master/LICENSE

package main

import (
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/maniksurtani/quotaservice"
	"github.com/maniksurtani/quotaservice/buckets/memory"
	"github.com/maniksurtani/quotaservice/config"
	"github.com/maniksurtani/quotaservice/rpc/grpc"
	"github.com/maniksurtani/quotaservice/stats"
)

func main() {
	cfg := config.NewDefaultServiceConfig()
	ns := config.NewDefaultNamespaceConfig("test.namespace")
	ns.DynamicBucketTemplate = config.NewDefaultBucketConfig(config.DynamicBucketTemplateName)
	ns.DynamicBucketTemplate.Size = 100000000000
	ns.DynamicBucketTemplate.FillRate = 100000000
	b := config.NewDefaultBucketConfig("xyz")
	config.AddBucket(ns, b)
	config.AddNamespace(cfg, ns)

	ns = config.NewDefaultNamespaceConfig("test.namespace2")
	ns.DefaultBucket = config.NewDefaultBucketConfig(config.DefaultBucketName)
	b = config.NewDefaultBucketConfig("xyz")
	config.AddBucket(ns, b)
	config.AddNamespace(cfg, ns)

	server := quotaservice.New(memory.NewBucketFactory(),
		config.NewMemoryConfigPersister(),
		cfg,
		grpc.New("localhost:10990"))
	server.SetStatsListener(stats.NewMemoryStatsListener())
	server.Start()

	// Serve Admin Console
	sm := http.NewServeMux()
	server.ServeAdminConsole(sm, "admin/public", true)
	go func() { http.ListenAndServe("localhost:8080", sm) }()

	// Block until SIGTERM, SIGKILL or SIGINT
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGTERM, syscall.SIGKILL, syscall.SIGINT)

	var shutdown sync.WaitGroup
	shutdown.Add(1)

	go func() {
		<-sigs
		shutdown.Done()
	}()

	shutdown.Wait()
	server.Stop()
}
<commit_msg>Log admin host and port when running test server<commit_after>// Licensed under the Apache License, Version 2.0
// Details: https://raw.githubusercontent.com/maniksurtani/quotaservice/master/LICENSE

package main

import (
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/maniksurtani/quotaservice"
	"github.com/maniksurtani/quotaservice/buckets/memory"
	"github.com/maniksurtani/quotaservice/config"
	"github.com/maniksurtani/quotaservice/logging"
	"github.com/maniksurtani/quotaservice/rpc/grpc"
	"github.com/maniksurtani/quotaservice/stats"
)

const (
	ADMIN_SERVER = "localhost:8080"
	GRPC_SERVER  = "localhost:10990"
)

func main() {
	cfg := config.NewDefaultServiceConfig()
	ns := config.NewDefaultNamespaceConfig("test.namespace")
	ns.DynamicBucketTemplate = config.NewDefaultBucketConfig(config.DynamicBucketTemplateName)
	ns.DynamicBucketTemplate.Size = 100000000000
	ns.DynamicBucketTemplate.FillRate = 100000000
	b := config.NewDefaultBucketConfig("xyz")
	config.AddBucket(ns, b)
	config.AddNamespace(cfg, ns)

	ns = config.NewDefaultNamespaceConfig("test.namespace2")
	ns.DefaultBucket = config.NewDefaultBucketConfig(config.DefaultBucketName)
	b = config.NewDefaultBucketConfig("xyz")
	config.AddBucket(ns, b)
	config.AddNamespace(cfg, ns)

	server := quotaservice.New(memory.NewBucketFactory(),
		config.NewMemoryConfigPersister(),
		cfg,
		grpc.New(GRPC_SERVER))
	server.SetStatsListener(stats.NewMemoryStatsListener())
	server.Start()

	// Serve Admin Console
	logging.Printf("Starting admin server on %v\n", ADMIN_SERVER)
	sm := http.NewServeMux()
	server.ServeAdminConsole(sm, "admin/public", true)
	go func() {
		http.ListenAndServe(ADMIN_SERVER, sm)
	}()

	// Block until SIGTERM, SIGKILL or SIGINT
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGTERM, syscall.SIGKILL, syscall.SIGINT)

	var shutdown sync.WaitGroup
	shutdown.Add(1)

	go func() {
		<-sigs
		shutdown.Done()
	}()

	shutdown.Wait()
	server.Stop()
}
<|endoftext|>


<commit_before>package driver

import (
	"reflect"
	"strconv"
	"testing"
	"time"
)

func testConcurrentVolatileKeyValueStore(t *testing.T, drv ConcurrentVolatileKeyValueStore) {
	defer drv.Close()

	expiDur := 10 * time.Millisecond

	// まだ無い。
	if v, _, err := drv.Get(testKey, nil); err != nil {
		t.Fatal(err)
	} else if v != nil {
		t.Error(v)
	}

	// 入れる。
	if _, err := drv.Put(testKey, testVal, time.Now().Add(time.Second)); err != nil {
		t.Fatal(err)
	}

	// ある。
	if v, _, err := drv.Get(testKey, nil); err != nil {
		t.Fatal(err)
	} else if !reflect.DeepEqual(v, testVal) {
		if !jsonEqual(v, testVal) {
			t.Error(v)
		}
	}

	// 消す。
	if err := drv.Remove(testKey); err != nil {
		t.Fatal(err)
	}

	// もう無い。
	if v, _, err := drv.Get(testKey, nil); err != nil {
		t.Fatal(err)
	} else if v != nil {
		t.Error(v)
	}

	// また入れる。
	if _, err := drv.Put(testKey, testVal, time.Now().Add(time.Second)); err != nil {
		t.Fatal(err)
	}

	// エントリが異なれば入れられない。
	exp := time.Now().Add(expiDur)
	if v, _, err := drv.GetAndSetEntry(testKey, nil, testKey+":entry", "0", exp); err != nil {
		t.Fatal(err)
	} else if !reflect.DeepEqual(v, testVal) {
		if !jsonEqual(v, testVal) {
			t.Error(v)
		}
	} else if ok, _, err := drv.PutIfEntered(testKey, testVal, exp, testKey+":entry", "1"); err != nil {
		t.Fatal(err)
	} else if ok {
		t.Fatal("")
	}

	bef := time.Now()
	if ok, _, err := drv.PutIfEntered(testKey, testVal, exp, testKey+":entry", "0"); err != nil {
		t.Fatal(err)
	} else if !ok {
		t.Fatal("")
	}
	diff := int64(time.Since(bef) / time.Nanosecond)

	// 消えるかどうか。
	for {
		bef := time.Now()
		v, _, err := drv.Get(testKey, nil)
		if err != nil {
			t.Fatal(err)
		}
		aft := time.Now()

		// GC 等で時間が掛かることもあるため、aft > exp でも nil が返るとは限らない。
		// だが、aft <= exp であれば非 nil が返らなければならない。
		// 同様に、bef > exp であれば nil が返らなければならない。

		if aft.UnixNano() <= cutOff(exp.UnixNano(), 1e6)-diff { // redis の粒度がミリ秒のため。
			if v == nil {
				t.Error(aft)
				t.Error(exp)
				return
			}
		} else if bef.UnixNano() > cutOff(exp.UnixNano(), 1e6)+1e6+diff { // redis の粒度がミリ秒のため。
			if v != nil {
				t.Error(bef)
				t.Error(exp)
				return
			}
			// 消えた。
			return
		} else if v == nil { // bef <= exp < aft
			// 消えた。
			return
		}

		time.Sleep(time.Millisecond)
	}
}

func testConcurrentVolatileKeyValueStoreConsistency(t *testing.T, drv ConcurrentVolatileKeyValueStore) {
	defer drv.Close()

	const n = 5
	const loop = 1000
	const expiDur = time.Minute

	if _, err := drv.Put(testKey, float64(0), time.Now().Add(expiDur)); err != nil {
		t.Fatal(err)
	}
	defer drv.Remove(testKey)

	const eKey = testKey + ":entry"
	defer drv.Remove(eKey)

	resCh := make(chan int, n)
	errCh := make(chan error, n)
	for i := 0; i < n; i++ {
		go func(id int) {
			eVal := strconv.Itoa(id)
			res := 0
			defer func() { resCh <- res }()
			for j := 0; j < loop; j++ {
				v, _, err := drv.GetAndSetEntry(testKey, nil, eKey, eVal, time.Now().Add(expiDur))
				if err != nil {
					errCh <- err
					return
				}
				a, _ := v.(float64)
				a += 1
				ok, _, err := drv.PutIfEntered(testKey, a, time.Now().Add(expiDur), eKey, eVal)
				if err != nil {
					errCh <- err
					return
				}
				if ok {
					res++
				}
			}
		}(i)
	}

	res := 0
	for i := 0; i < n; i++ {
		res += <-resCh
	}

	v, _, err := drv.Get(testKey, nil)
	if err != nil {
		errCh <- err
		return
	}
	a, _ := v.(float64)
	if int(a) != res {
		t.Error(a, res)
	}

	for ok := false; !ok; {
		select {
		case err := <-errCh:
			t.Error(err)
		default:
			ok = true
		}
	}
}
<commit_msg>低性能マシン向けテストパラメータ調整<commit_after>package driver

import (
	"reflect"
	"strconv"
	"testing"
	"time"
)

func testConcurrentVolatileKeyValueStore(t *testing.T, drv ConcurrentVolatileKeyValueStore) {
	defer drv.Close()

	expiDur := 10 * time.Millisecond

	// まだ無い。
	if v, _, err := drv.Get(testKey, nil); err != nil {
		t.Fatal(err)
	} else if v != nil {
		t.Error(v)
	}

	// 入れる。
	if _, err := drv.Put(testKey, testVal, time.Now().Add(time.Second)); err != nil {
		t.Fatal(err)
	}

	// ある。
	if v, _, err := drv.Get(testKey, nil); err != nil {
		t.Fatal(err)
	} else if !reflect.DeepEqual(v, testVal) {
		if !jsonEqual(v, testVal) {
			t.Error(v)
		}
	}

	// 消す。
	if err := drv.Remove(testKey); err != nil {
		t.Fatal(err)
	}

	// もう無い。
	if v, _, err := drv.Get(testKey, nil); err != nil {
		t.Fatal(err)
	} else if v != nil {
		t.Error(v)
	}

	// また入れる。
	if _, err := drv.Put(testKey, testVal, time.Now().Add(time.Second)); err != nil {
		t.Fatal(err)
	}

	// エントリが異なれば入れられない。
	if v, _, err := drv.GetAndSetEntry(testKey, nil, testKey+":entry", "0", time.Now().Add(time.Second)); err != nil {
		t.Fatal(err)
	} else if !reflect.DeepEqual(v, testVal) {
		if !jsonEqual(v, testVal) {
			t.Error(v)
		}
	}

	exp := time.Now().Add(expiDur)
	if ok, _, err := drv.PutIfEntered(testKey, testVal, exp, testKey+":entry", "1"); err != nil {
		t.Fatal(err)
	} else if ok {
		t.Fatal("")
	}

	bef := time.Now()
	if ok, _, err := drv.PutIfEntered(testKey, testVal, exp, testKey+":entry", "0"); err != nil {
		t.Fatal(err)
	} else if !ok {
		t.Fatal("")
	}
	diff := int64(time.Since(bef) / time.Nanosecond)

	// 消えるかどうか。
	for {
		bef := time.Now()
		v, _, err := drv.Get(testKey, nil)
		if err != nil {
			t.Fatal(err)
		}
		aft := time.Now()

		// GC 等で時間が掛かることもあるため、aft > exp でも nil が返るとは限らない。
		// だが、aft <= exp であれば非 nil が返らなければならない。
		// 同様に、bef > exp であれば nil が返らなければならない。

		if aft.UnixNano() <= cutOff(exp.UnixNano(), 1e6)-diff { // redis の粒度がミリ秒のため。
			if v == nil {
				t.Error(aft)
				t.Error(exp)
				return
			}
		} else if bef.UnixNano() > cutOff(exp.UnixNano(), 1e6)+1e6+diff { // redis の粒度がミリ秒のため。
			if v != nil {
				t.Error(bef)
				t.Error(exp)
				return
			}
			// 消えた。
			return
		} else if v == nil { // bef <= exp < aft
			// 消えた。
			return
		}

		time.Sleep(time.Millisecond)
	}
}

func testConcurrentVolatileKeyValueStoreConsistency(t *testing.T, drv ConcurrentVolatileKeyValueStore) {
	defer drv.Close()

	const n = 5
	const loop = 1000
	const expiDur = time.Minute

	if _, err := drv.Put(testKey, float64(0), time.Now().Add(expiDur)); err != nil {
		t.Fatal(err)
	}
	defer drv.Remove(testKey)

	const eKey = testKey + ":entry"
	defer drv.Remove(eKey)

	resCh := make(chan int, n)
	errCh := make(chan error, n)
	for i := 0; i < n; i++ {
		go func(id int) {
			eVal := strconv.Itoa(id)
			res := 0
			defer func() { resCh <- res }()
			for j := 0; j < loop; j++ {
				v, _, err := drv.GetAndSetEntry(testKey, nil, eKey, eVal, time.Now().Add(expiDur))
				if err != nil {
					errCh <- err
					return
				}
				a, _ := v.(float64)
				a += 1
				ok, _, err := drv.PutIfEntered(testKey, a, time.Now().Add(expiDur), eKey, eVal)
				if err != nil {
					errCh <- err
					return
				}
				if ok {
					res++
				}
			}
		}(i)
	}

	res := 0
	for i := 0; i < n; i++ {
		res += <-resCh
	}

	v, _, err := drv.Get(testKey, nil)
	if err != nil {
		errCh <- err
		return
	}
	a, _ := v.(float64)
	if int(a) != res {
		t.Error(a, res)
	}

	for ok := false; !ok; {
		select {
		case err := <-errCh:
			t.Error(err)
		default:
			ok = true
		}
	}
}
<|endoftext|>


<commit_before>package web

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"fmt"
	"math"
	"net/http"
	"net/url"
	"sort"
	"sync"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/getlantern/zenodb/core"
	"github.com/getlantern/zenodb/encoding"
	"github.com/getlantern/zenodb/sql"
	"github.com/gorilla/mux"
	"github.com/retailnext/hllpp"
)

const (
	nanosPerMilli = 1000000

	pauseTime    = 250 * time.Millisecond
	shortTimeout = 5 * time.Second
	longTimeout  = 1000 * time.Hour
)

type QueryResult struct {
	SQL                string
	Permalink          string
	TS                 int64
	TSCardinality      uint64
	Fields             []string
	FieldCardinalities []uint64
	Dims               []string
	DimCardinalities   []uint64
	Rows               []*ResultRow
}

type ResultRow struct {
	TS   int64
	Key  map[string]interface{}
	Vals []float64
}

type query struct {
	sqlString string
	parsed    *sql.Query
	ce        cacheEntry
}

func (h *handler) runQuery(resp http.ResponseWriter, req *http.Request) {
	h.sqlQuery(resp, req, longTimeout)
}

func (h *handler) asyncQuery(resp http.ResponseWriter, req *http.Request) {
	h.sqlQuery(resp, req, shortTimeout)
}

func (h *handler) cachedQuery(resp http.ResponseWriter, req *http.Request) {
	if !h.authenticate(resp, req) {
		resp.WriteHeader(http.StatusForbidden)
		return
	}

	log.Debug(req.URL)
	permalink := mux.Vars(req)["permalink"]
	ce, err := h.cache.getByPermalink(permalink)
	if ce == nil {
		http.NotFound(resp, req)
		return
	}
	h.respondWithCacheEntry(resp, req, ce, err, shortTimeout)
}

func (h *handler) sqlQuery(resp http.ResponseWriter, req *http.Request, timeout time.Duration) {
	if !h.authenticate(resp, req) {
		resp.WriteHeader(http.StatusForbidden)
		return
	}

	log.Debug(req.URL)
	sqlString, _ := url.QueryUnescape(req.URL.RawQuery)

	ce, err := h.query(req, sqlString)
	h.respondWithCacheEntry(resp, req, ce, err, timeout)
}

func (h *handler) respondWithCacheEntry(resp http.ResponseWriter, req *http.Request, ce cacheEntry, err error, timeout time.Duration) {
	limit := int(timeout / pauseTime)
	for i := 0; i < limit; i++ {
		if err != nil {
			log.Error(err)
			resp.WriteHeader(http.StatusInternalServerError)
			fmt.Fprint(resp, err.Error())
			return
		}
		switch ce.status() {
		case statusSuccess:
			h.respondSuccess(resp, req, ce)
			return
		case statusError:
			h.respondError(resp, req, ce)
			return
		case statusPending:
			// Pause a little bit and try again
			time.Sleep(pauseTime)
			ce, err = h.cache.getByPermalink(ce.permalink())
		}
	}
	// Let the client know that we're still working on it
	resp.WriteHeader(http.StatusAccepted)
	fmt.Fprintf(resp, "/cached/%v", ce.permalink())
}

func (h *handler) respondSuccess(resp http.ResponseWriter, req *http.Request, ce cacheEntry) {
	resp.Header().Set("Content-Type", "application/json")
	resp.Header().Set("Expires", "0")
	resp.Header().Set("Cache-control", "no-cache, no-store, must-revalidate")
	resp.Header().Set("Content-Encoding", "gzip")
	resp.WriteHeader(http.StatusOK)
	resp.Write(ce.data())
}

func (h *handler) respondError(resp http.ResponseWriter, req *http.Request, ce cacheEntry) {
	resp.WriteHeader(http.StatusInternalServerError)
	resp.Write(ce.error())
}

func (h *handler) query(req *http.Request, sqlString string) (ce cacheEntry, err error) {
	parsed, parseErr := sql.Parse(sqlString)
	if parseErr != nil {
		return nil, parseErr
	}

	if req.Header.Get("Cache-control") == "no-cache" {
		ce, err = h.cache.begin(sqlString)
		if err != nil {
			return
		}
	} else {
		var created bool
		ce, created, err = h.cache.getOrBegin(sqlString)
		if err != nil || !created {
			return
		}
		if ce.status() != statusPending {
			log.Debugf("Found results for %v in cache", sqlString)
			return
		}
	}

	// Request query to run in background
	h.queries <- &query{sqlString, parsed, ce}

	return
}

func (h *handler) coalesceQueries() {
	for {
		var coalescedQueries []*query
		var remainingQueries []*query
		table := ""
	coalesceLoop:
		for {
			select {
			case query := <-h.queries:
				if table == "" || table == query.parsed.From {
					coalescedQueries = append(coalescedQueries)
					table = query.parsed.From
				} else {
					remainingQueries = append(remainingQueries, query)
				}
			case <-time.After(15 * time.Second):
				break coalesceLoop
			}
		}

		log.Debugf("Coalesced %d queries to %v", len(coalescedQueries), table)
		// re-queue queries that weren't included in this run
		for _, query := range remainingQueries {
			h.queries <- query
		}
		h.coalescedQueries <- coalescedQueries
	}
}

func (h *handler) processQueries() {
	for queries := range h.coalescedQueries {
		for _, query := range queries {
			sqlString := query.sqlString
			ce := query.ce
			result, err := h.doQuery(sqlString, ce.permalink())
			if err != nil {
				err = fmt.Errorf("Unable to query: %v", err)
				log.Error(err)
				ce = ce.fail(err)
			} else {
				resultBytes, err := compress(json.Marshal(result))
				if err != nil {
					err = fmt.Errorf("Unable to marshal result: %v", err)
					log.Error(err)
					ce = ce.fail(err)
				} else if len(resultBytes) > h.MaxResponseBytes {
					err = fmt.Errorf("Query result size %v exceeded limit of %v", humanize.Bytes(uint64(len(resultBytes))), humanize.Bytes(uint64(h.MaxResponseBytes)))
					log.Error(err)
					ce = ce.fail(err)
				} else {
					ce = ce.succeed(resultBytes)
				}
			}
			h.cache.put(sqlString, ce)
			log.Debugf("Cached results for %v", sqlString)
		}
	}
}

func compress(resultBytes []byte, err error) ([]byte, error) {
	if err != nil {
		return nil, err
	}
	buf := bytes.NewBuffer(make([]byte, 0, len(resultBytes)))
	gw, err := gzip.NewWriterLevel(buf, gzip.BestCompression)
	if err != nil {
		return nil, err
	}
	_, err = gw.Write(resultBytes)
	if err != nil {
		return nil, err
	}
	err = gw.Close()
	if err != nil {
		return nil, err
	}
	compressed := buf.Bytes()
	log.Debugf("Compressed result from %v down to %v using gzip", humanize.Bytes(uint64(len(resultBytes))), humanize.Bytes(uint64(len(compressed))))
	return compressed, nil
}

func (h *handler) doQuery(sqlString string, permalink string) (*QueryResult, error) {
	rs, err := h.db.Query(sqlString, false, nil, false)
	if err != nil {
		return nil, err
	}

	var addDim func(dim string)

	result := &QueryResult{
		SQL:       sqlString,
		Permalink: permalink,
		TS:        time.Now().UnixNano() / nanosPerMilli,
	}
	groupBy := rs.GetGroupBy()
	if len(groupBy) > 0 {
		addDim = func(dim string) {
			// noop
		}
		for _, gb := range groupBy {
			result.Dims = append(result.Dims, gb.Name)
		}
	} else {
		addDim = func(dim string) {
			found := false
			for _, existing := range result.Dims {
				if existing == dim {
					found = true
					break
				}
			}
			if !found {
				result.Dims = append(result.Dims, dim)
			}
		}
	}

	var fields core.Fields
	var fieldCardinalities []*hllpp.HLLPP
	dimCardinalities := make(map[string]*hllpp.HLLPP)
	tsCardinality := hllpp.New()
	cbytes := make([]byte, 8)

	estimatedResultBytes := 0
	var mx sync.Mutex
	ctx, cancel := context.WithTimeout(context.Background(), h.QueryTimeout)
	defer cancel()
	rs.Iterate(ctx, func(inFields core.Fields) error {
		fields = inFields
		for _, field := range fields {
			result.Fields = append(result.Fields, field.Name)
			fieldCardinalities = append(fieldCardinalities, hllpp.New())
		}
		return nil
	}, func(row *core.FlatRow) (bool, error) {
		mx.Lock()
		key := make(map[string]interface{}, 10)
		row.Key.Iterate(true, true, func(dim string, value interface{}, valueBytes []byte) bool {
			key[dim] = value
			addDim(dim)
			hlp := dimCardinalities[dim]
			if hlp == nil {
				hlp = hllpp.New()
				dimCardinalities[dim] = hlp
			}
			hlp.Add(valueBytes)
			estimatedResultBytes += len(dim) + len(valueBytes)
			return true
		})

		estimatedResultBytes += 8 * len(row.Values)
		if estimatedResultBytes > h.MaxResponseBytes {
			mx.Unlock()
			// Note - the estimated size here is always an underestimate of the final
			// JSON size, so this is a conservative way to check. The final check
			// after generating the JSON may sometimes catch things that slipped
			// through here.
			return false, fmt.Errorf("Estimated query result size %v exceeded limit of %v", humanize.Bytes(uint64(estimatedResultBytes)), humanize.Bytes(uint64(h.MaxResponseBytes)))
		}

		encoding.Binary.PutUint64(cbytes, uint64(row.TS))
		tsCardinality.Add(cbytes)

		resultRow := &ResultRow{
			TS:   row.TS / nanosPerMilli,
			Key:  key,
			Vals: make([]float64, 0, len(row.Values)),
		}

		for i, value := range row.Values {
			resultRow.Vals = append(resultRow.Vals, value)
			encoding.Binary.PutUint64(cbytes, math.Float64bits(value))
			fieldCardinalities[i].Add(cbytes)
		}
		result.Rows = append(result.Rows, resultRow)
		mx.Unlock()
		return true, nil
	})

	result.TSCardinality = tsCardinality.Count()
	result.Dims = make([]string, 0, len(dimCardinalities))
	for dim := range dimCardinalities {
		result.Dims = append(result.Dims, dim)
	}
	sort.Strings(result.Dims)
	for _, dim := range result.Dims {
		result.DimCardinalities = append(result.DimCardinalities, dimCardinalities[dim].Count())
	}

	result.FieldCardinalities = make([]uint64, 0, len(fieldCardinalities))
	for _, fieldCardinality := range fieldCardinalities {
		result.FieldCardinalities = append(result.FieldCardinalities, fieldCardinality.Count())
	}

	return result, nil
}

func intToBytes(i uint64) []byte {
	b := make([]byte, 8)
	encoding.Binary.PutUint64(b, i)
	return b
}
<commit_msg>Cleaned up web coalescing<commit_after>package web

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"fmt"
	"math"
	"net/http"
	"net/url"
	"sort"
	"sync"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/getlantern/zenodb/core"
	"github.com/getlantern/zenodb/encoding"
	"github.com/getlantern/zenodb/sql"
	"github.com/gorilla/mux"
	"github.com/retailnext/hllpp"
)

const (
	nanosPerMilli = 1000000

	pauseTime    = 250 * time.Millisecond
	shortTimeout = 5 * time.Second
	longTimeout  = 1000 * time.Hour
)

type QueryResult struct {
	SQL                string
	Permalink          string
	TS                 int64
	TSCardinality      uint64
	Fields             []string
	FieldCardinalities []uint64
	Dims               []string
	DimCardinalities   []uint64
	Rows               []*ResultRow
}

type ResultRow struct {
	TS   int64
	Key  map[string]interface{}
	Vals []float64
}

type query struct {
	sqlString string
	parsed    *sql.Query
	ce        cacheEntry
}

func (h *handler) runQuery(resp http.ResponseWriter, req *http.Request) {
	h.sqlQuery(resp, req, longTimeout)
}

func (h *handler) asyncQuery(resp http.ResponseWriter, req *http.Request) {
	h.sqlQuery(resp, req, shortTimeout)
}

func (h *handler) cachedQuery(resp http.ResponseWriter, req *http.Request) {
	if !h.authenticate(resp, req) {
		resp.WriteHeader(http.StatusForbidden)
		return
	}

	log.Debug(req.URL)
	permalink := mux.Vars(req)["permalink"]
	ce, err := h.cache.getByPermalink(permalink)
	if ce == nil {
		http.NotFound(resp, req)
		return
	}
	h.respondWithCacheEntry(resp, req, ce, err, shortTimeout)
}

func (h *handler) sqlQuery(resp http.ResponseWriter, req *http.Request, timeout time.Duration) {
	if !h.authenticate(resp, req) {
		resp.WriteHeader(http.StatusForbidden)
		return
	}

	log.Debug(req.URL)
	sqlString, _ := url.QueryUnescape(req.URL.RawQuery)

	ce, err := h.query(req, sqlString)
	h.respondWithCacheEntry(resp, req, ce, err, timeout)
}

func (h *handler) respondWithCacheEntry(resp http.ResponseWriter, req *http.Request, ce cacheEntry, err error, timeout time.Duration) {
	limit := int(timeout / pauseTime)
	for i := 0; i < limit; i++ {
		if err != nil {
			log.Error(err)
			resp.WriteHeader(http.StatusInternalServerError)
			fmt.Fprint(resp, err.Error())
			return
		}
		switch ce.status() {
		case statusSuccess:
			h.respondSuccess(resp, req, ce)
			return
		case statusError:
			h.respondError(resp, req, ce)
			return
		case statusPending:
			// Pause a little bit and try again
			time.Sleep(pauseTime)
			ce, err = h.cache.getByPermalink(ce.permalink())
		}
	}
	// Let the client know that we're still working on it
	resp.WriteHeader(http.StatusAccepted)
	fmt.Fprintf(resp, "/cached/%v", ce.permalink())
}

func (h *handler) respondSuccess(resp http.ResponseWriter, req *http.Request, ce cacheEntry) {
	resp.Header().Set("Content-Type", "application/json")
	resp.Header().Set("Expires", "0")
	resp.Header().Set("Cache-control", "no-cache, no-store, must-revalidate")
	resp.Header().Set("Content-Encoding", "gzip")
	resp.WriteHeader(http.StatusOK)
	resp.Write(ce.data())
}

func (h *handler) respondError(resp http.ResponseWriter, req *http.Request, ce cacheEntry) {
	resp.WriteHeader(http.StatusInternalServerError)
	resp.Write(ce.error())
}

func (h *handler) query(req *http.Request, sqlString string) (ce cacheEntry, err error) {
	parsed, parseErr := sql.Parse(sqlString)
	if parseErr != nil {
		return nil, parseErr
	}

	if req.Header.Get("Cache-control") == "no-cache" {
		ce, err = h.cache.begin(sqlString)
		if err != nil {
			return
		}
	} else {
		var created bool
		ce, created, err = h.cache.getOrBegin(sqlString)
		if err != nil || !created {
			return
		}
		if ce.status() != statusPending {
			log.Debugf("Found results for %v in cache", sqlString)
			return
		}
	}

	// Request query to run in background
	h.queries <- &query{sqlString, parsed, ce}

	return
}

func (h *handler) coalesceQueries() {
	for q := range h.queries {
		coalescedQueries := append([]*query(nil), q)
		var remainingQueries []*query
		table := q.parsed.From
	coalesceLoop:
		for {
			select {
			case query := <-h.queries:
				if table == query.parsed.From {
					coalescedQueries = append(coalescedQueries)
				} else {
					remainingQueries = append(remainingQueries, query)
				}
			case <-time.After(15 * time.Second):
				break coalesceLoop
			}
		}

		log.Debugf("Coalesced %d queries to %v", len(coalescedQueries), table)
		// re-queue queries that weren't included in this run
		for _, query := range remainingQueries {
			h.queries <- query
		}
		h.coalescedQueries <- coalescedQueries
	}
}

func (h *handler) processQueries() {
	for queries := range h.coalescedQueries {
		for _, query := range queries {
			sqlString := query.sqlString
			ce := query.ce
			result, err := h.doQuery(sqlString, ce.permalink())
			if err != nil {
				err = fmt.Errorf("Unable to query: %v", err)
				log.Error(err)
				ce = ce.fail(err)
			} else {
				resultBytes, err := compress(json.Marshal(result))
				if err != nil {
					err = fmt.Errorf("Unable to marshal result: %v", err)
					log.Error(err)
					ce = ce.fail(err)
				} else if len(resultBytes) > h.MaxResponseBytes {
					err = fmt.Errorf("Query result size %v exceeded limit of %v", humanize.Bytes(uint64(len(resultBytes))), humanize.Bytes(uint64(h.MaxResponseBytes)))
					log.Error(err)
					ce = ce.fail(err)
				} else {
					ce = ce.succeed(resultBytes)
				}
			}
			h.cache.put(sqlString, ce)
			log.Debugf("Cached results for %v", sqlString)
		}
	}
}

func compress(resultBytes []byte, err error) ([]byte, error) {
	if err != nil {
		return nil, err
	}
	buf := bytes.NewBuffer(make([]byte, 0, len(resultBytes)))
	gw, err := gzip.NewWriterLevel(buf, gzip.BestCompression)
	if err != nil {
		return nil, err
	}
	_, err = gw.Write(resultBytes)
	if err != nil {
		return nil, err
	}
	err = gw.Close()
	if err != nil {
		return nil, err
	}
	compressed := buf.Bytes()
	log.Debugf("Compressed result from %v down to %v using gzip", humanize.Bytes(uint64(len(resultBytes))), humanize.Bytes(uint64(len(compressed))))
	return compressed, nil
}

func (h *handler) doQuery(sqlString string, permalink string) (*QueryResult, error) {
	rs, err := h.db.Query(sqlString, false, nil, false)
	if err != nil {
		return nil, err
	}

	var addDim func(dim string)

	result := &QueryResult{
		SQL:       sqlString,
		Permalink: permalink,
		TS:        time.Now().UnixNano() / nanosPerMilli,
	}
	groupBy := rs.GetGroupBy()
	if len(groupBy) > 0 {
		addDim = func(dim string) {
			// noop
		}
		for _, gb := range groupBy {
			result.Dims = append(result.Dims, gb.Name)
		}
	} else {
		addDim = func(dim string) {
			found := false
			for _, existing := range result.Dims {
				if existing == dim {
					found = true
					break
				}
			}
			if !found {
				result.Dims = append(result.Dims, dim)
			}
		}
	}

	var fields core.Fields
	var fieldCardinalities []*hllpp.HLLPP
	dimCardinalities := make(map[string]*hllpp.HLLPP)
	tsCardinality := hllpp.New()
	cbytes := make([]byte, 8)

	estimatedResultBytes := 0
	var mx sync.Mutex
	ctx, cancel := context.WithTimeout(context.Background(), h.QueryTimeout)
	defer cancel()
	rs.Iterate(ctx, func(inFields core.Fields) error {
		fields = inFields
		for _, field := range fields {
			result.Fields = append(result.Fields, field.Name)
			fieldCardinalities = append(fieldCardinalities, hllpp.New())
		}
		return nil
	}, func(row *core.FlatRow) (bool, error) {
		mx.Lock()
		key := make(map[string]interface{}, 10)
		row.Key.Iterate(true, true, func(dim string, value interface{}, valueBytes []byte) bool {
			key[dim] = value
			addDim(dim)
			hlp := dimCardinalities[dim]
			if hlp == nil {
				hlp = hllpp.New()
				dimCardinalities[dim] = hlp
			}
			hlp.Add(valueBytes)
			estimatedResultBytes += len(dim) + len(valueBytes)
			return true
		})

		estimatedResultBytes += 8 * len(row.Values)
		if estimatedResultBytes > h.MaxResponseBytes {
			mx.Unlock()
			// Note - the estimated size here is always an underestimate of the final
			// JSON size, so this is a conservative way to check. The final check
			// after generating the JSON may sometimes catch things that slipped
			// through here.
			return false, fmt.Errorf("Estimated query result size %v exceeded limit of %v", humanize.Bytes(uint64(estimatedResultBytes)), humanize.Bytes(uint64(h.MaxResponseBytes)))
		}

		encoding.Binary.PutUint64(cbytes, uint64(row.TS))
		tsCardinality.Add(cbytes)

		resultRow := &ResultRow{
			TS:   row.TS / nanosPerMilli,
			Key:  key,
			Vals: make([]float64, 0, len(row.Values)),
		}

		for i, value := range row.Values {
			resultRow.Vals = append(resultRow.Vals, value)
			encoding.Binary.PutUint64(cbytes, math.Float64bits(value))
			fieldCardinalities[i].Add(cbytes)
		}
		result.Rows = append(result.Rows, resultRow)
		mx.Unlock()
		return true, nil
	})

	result.TSCardinality = tsCardinality.Count()
	result.Dims = make([]string, 0, len(dimCardinalities))
	for dim := range dimCardinalities {
		result.Dims = append(result.Dims, dim)
	}
	sort.Strings(result.Dims)
	for _, dim := range result.Dims {
		result.DimCardinalities = append(result.DimCardinalities, dimCardinalities[dim].Count())
	}

	result.FieldCardinalities = make([]uint64, 0, len(fieldCardinalities))
	for _, fieldCardinality := range fieldCardinalities {
		result.FieldCardinalities = append(result.FieldCardinalities, fieldCardinality.Count())
	}

	return result, nil
}

func intToBytes(i uint64) []byte {
	b := make([]byte, 8)
	encoding.Binary.PutUint64(b, i)
	return b
}
<|endoftext|>


<commit_before>package wiki

import (
	"encoding/json"
	"os"

	httpdate "github.com/Songmu/go-httpdate"
	"github.com/cooper/quiki/wikifier"
)

// DisplayPage represents a page result to display.
type DisplayPage struct {

	// basename of the page, with the extension
	File string `json:"file,omitempty"`

	// basename of the page, without the extension
	Name string `json:"name,omitempty"`

	// absolute file path of the page
	Path string `json:"path,omitempty"`

	// the page content (HTML)
	Content wikifier.HTML `json:"content,omitempty"`

	// UNIX timestamp of when the page was last modified.
	// if Generated is true, this is the current time.
	// if FromCache is true, this is the modified date of the cache file.
	// otherwise, this is the modified date of the page file itself.
	ModUnix int64 `json:"mod_unix,omitempty"`

	// like ModUnix except in HTTP date format, suitable for Last-Modified
	Modified string `json:"modified,omitempty"`

	// CSS generated for the page from style{} blocks
	CSS string `json:"css,omitempty"`

	// true if this content was read from a cache file. opposite of Generated
	FromCache bool `json:"cached,omitempty"`

	// true if the content being served was just generated on the fly.
	// opposite of FromCache
	Generated bool `json:"generated,omitempty"`

	// true if this request resulted in the writing of a new cache file.
	// this can only be true if Generated is true
	CacheGenerated bool `json:"cache_gen,omitempty"`

	// true if this request resulted in the writing of a text file.
	// this can only be true if Generated is true
	TextGenerated bool `json:"text_gen,omitempty"`

	// true if the page has not yet been published for public viewing.
	// this only occurs when it is specified that serving drafts is OK,
	// since normally a draft page instead results in a DisplayError.
	Draft bool `json:"draft,omitempty"`

	// warnings produced by the parser
	Warnings []string `json:"warnings,omitempty"`

	// UNIX timestamp of when the page was created, as extracted from
	// the special @page.created variable
	CreatedUnix int64 `json:"created,omitempty"`

	// name of the page author, as extracted from the special @page.author
	// variable
	Author string `json:"author,omitempty"`

	// list of categories the page belongs to, without the '.cat' extension
	Categories []string `json:"categories,omitempty"`

	// page title as extracted from the special @page.title variable, including
	// any possible HTML-encoded formatting
	FmtTitle wikifier.HTML `json:"fmt_title,omitempty"`

	// like FmtTitle except that all text formatting has been stripped.
	// suitable for use in the <title> tag
	Title string `json:"title,omitempty"`
}

// NewPage creates a Page given its filepath and configures it for
// use with this Wiki.
func (w *Wiki) NewPage(name string) *wikifier.Page {
	p := wikifier.NewPage(w.pathForPage(name, false, ""))
	p.Opt = &w.Opt
	return p
}

// DisplayPage returns the display result for a page.
func (w *Wiki) DisplayPage(name string) interface{} {
	return w.DisplayPageDraft(name, false)
}

// DisplayPageDraft returns the display result for a page.
//
// Unlike DisplayPage, if draftOK is true, the content is served even if it is
// marked as draft.
//
func (w *Wiki) DisplayPageDraft(name string, draftOK bool) interface{} {
	var r DisplayPage

	// create the page
	page := w.NewPage(name)

	// file does not exist
	if !page.Exists() {
		return DisplayError{
			Error:         "Page does not exist.",
			DetailedError: "Page '" + page.FilePath + "' does not exist.",
		}
	}

	// filename and path info
	path := page.Path()
	r.File = page.Name()
	r.Name = page.NameNE()
	r.Path = path

	// FIRST redirect check - symbolic link
	if page.IsSymlink() {
		return DisplayRedirect{Redirect: page.Redirect()}
	}

	// caching is enabled, so serve the cached copy if available
	if w.Opt.Page.EnableCache && page.CacheExists() {
		w.displayCachedPage(page, &r)
		return r
	}

	// Safe point - we will be generating the page right now.

	// parse the page
	//
	// if an error occurs, parse it again in variable-only mode.
	// then hopefully we can at least get the metadata and categories
	//
	err := page.Parse()
	if err != nil {
		page.VarsOnly = true
		page.Parse()
		// TODO: add page to categories
		// TODO: cache the error
		return DisplayError{Error: err.Error(), ParseError: true}
	}

	// if this is a draft and we're not serving drafts, pretend
	// that the page does not exist
	if !draftOK && page.Draft() {
		return DisplayError{Error: "Page has not yet been publised.", Draft: true}
	}

	// SECOND redirect check -
	// this is for pages we just parsed with @page.redirect
	if redir := page.Redirect(); redir != "" {
		return DisplayRedirect{Redirect: redir}
	}

	// generate HTML and metadata
	r.Generated = true
	r.Title = page.Title()
	r.FmtTitle = page.FmtTitle()
	r.Author = page.Author()
	r.Draft = page.Draft()
	r.ModUnix = page.Modified().Unix()
	r.Modified = httpdate.Time2Str(page.Modified())
	r.Content = page.HTML()
	r.CSS = page.CSS()

	// TODO: should we include the page object?
	// TODO: warnings

	// TODO: update categories and set to r.Categories

	// write cache file if enabled
	if err := w.writePageCache(page, &r); err != nil {
		return *err
	}

	// TODO: write search file if enabled

	return r
}

func (w *Wiki) writePageCache(page *wikifier.Page, r *DisplayPage) *DisplayError {

	// caching isn't enabled
	if !page.Opt.Page.EnableCache || page.CachePath() == "" {
		return nil
	}

	// open the cache file for writing
	cacheFile, err := os.Create(page.CachePath())
	defer cacheFile.Close()
	if err != nil {
		return &DisplayError{
			Error:         "Could not write page cache file.",
			DetailedError: "Open '" + page.CachePath() + "' error: " + err.Error(),
		}
	}

	// generate page info
	j, err := json.Marshal(page.Info())
	if err != nil {
		return &DisplayError{
			Error:         "Could not write page cache file.",
			DetailedError: "JSON encode error: " + err.Error(),
		}
	}

	// save prefixing data
	cacheFile.Write(j)
	cacheFile.Write([]byte{'\n'})

	// save content
	content := string(r.Content)
	cacheFile.WriteString(content)
	if len(content) != 0 && content[len(content)-1] != '\n' {
		cacheFile.Write([]byte{'\n'})
	}

	// update result with real cache modified times
	r.ModUnix = page.CacheModified().Unix()
	r.Modified = httpdate.Time2Str(page.CacheModified())
	r.CacheGenerated = true

	return nil // success
}

func (w *Wiki) displayCachedPage(page *wikifier.Page, r *DisplayPage) {
}
<commit_msg>read page cache<commit_after>package wiki

import (
	"bytes"
	"encoding/json"
	"io/ioutil"
	"os"

	httpdate "github.com/Songmu/go-httpdate"
	"github.com/cooper/quiki/wikifier"
)

// DisplayPage represents a page result to display.
type DisplayPage struct {

	// basename of the page, with the extension
	File string `json:"file,omitempty"`

	// basename of the page, without the extension
	Name string `json:"name,omitempty"`

	// absolute file path of the page
	Path string `json:"path,omitempty"`

	// the page content (HTML)
	Content wikifier.HTML `json:"content,omitempty"`

	// UNIX timestamp of when the page was last modified.
	// if Generated is true, this is the current time.
	// if FromCache is true, this is the modified date of the cache file.
	// otherwise, this is the modified date of the page file itself.
	ModUnix int64 `json:"mod_unix,omitempty"`

	// like ModUnix except in HTTP date format, suitable for Last-Modified
	Modified string `json:"modified,omitempty"`

	// CSS generated for the page from style{} blocks
	CSS string `json:"css,omitempty"`

	// true if this content was read from a cache file. opposite of Generated
	FromCache bool `json:"cached,omitempty"`

	// true if the content being served was just generated on the fly.
	// opposite of FromCache
	Generated bool `json:"generated,omitempty"`

	// true if this request resulted in the writing of a new cache file.
	// this can only be true if Generated is true
	CacheGenerated bool `json:"cache_gen,omitempty"`

	// true if this request resulted in the writing of a text file.
	// this can only be true if Generated is true
	TextGenerated bool `json:"text_gen,omitempty"`

	// true if the page has not yet been published for public viewing.
	// this only occurs when it is specified that serving drafts is OK,
	// since normally a draft page instead results in a DisplayError.
	Draft bool `json:"draft,omitempty"`

	// warnings produced by the parser
	Warnings []string `json:"warnings,omitempty"`

	// UNIX timestamp of when the page was created, as extracted from
	// the special @page.created variable
	CreatedUnix int64 `json:"created,omitempty"`

	// name of the page author, as extracted from the special @page.author
	// variable
	Author string `json:"author,omitempty"`

	// list of categories the page belongs to, without the '.cat' extension
	Categories []string `json:"categories,omitempty"`

	// page title as extracted from the special @page.title variable, including
	// any possible HTML-encoded formatting
	FmtTitle wikifier.HTML `json:"fmt_title,omitempty"`

	// like FmtTitle except that all text formatting has been stripped.
	// suitable for use in the <title> tag
	Title string `json:"title,omitempty"`
}

type pageJSONManifest struct {
	CSS        string   `json:"css,omitempty"`
	Categories []string `json:"categories,omitempty"`
	Warnings   []string `json:"warnings,omitempty"`
	Error      string   `json:"error,omitempty"`
	wikifier.PageInfo
}

// NewPage creates a Page given its filepath and configures it for
// use with this Wiki.
func (w *Wiki) NewPage(name string) *wikifier.Page {
	p := wikifier.NewPage(w.pathForPage(name, false, ""))
	p.Opt = &w.Opt
	return p
}

// DisplayPage returns the display result for a page.
func (w *Wiki) DisplayPage(name string) interface{} {
	return w.DisplayPageDraft(name, false)
}

// DisplayPageDraft returns the display result for a page.
//
// Unlike DisplayPage, if draftOK is true, the content is served even if it is
// marked as draft.
//
func (w *Wiki) DisplayPageDraft(name string, draftOK bool) interface{} {
	var r DisplayPage

	// create the page
	page := w.NewPage(name)

	// file does not exist
	if !page.Exists() {
		return DisplayError{
			Error:         "Page does not exist.",
			DetailedError: "Page '" + page.FilePath + "' does not exist.",
		}
	}

	// filename and path info
	path := page.Path()
	r.File = page.Name()
	r.Name = page.NameNE()
	r.Path = path

	// FIRST redirect check - symbolic link
	if page.IsSymlink() {
		return DisplayRedirect{Redirect: page.Redirect()}
	}

	// caching is enabled, so serve the cached copy if available
	if w.Opt.Page.EnableCache && page.CacheExists() {
		if errOrRedir := w.displayCachedPage(page, &r, draftOK); errOrRedir != nil {
			return errOrRedir
		}
		return r
	}

	// Safe point - we will be generating the page right now.

	// parse the page
	//
	// if an error occurs, parse it again in variable-only mode.
	// then hopefully we can at least get the metadata and categories
	//
	err := page.Parse()
	if err != nil {
		page.VarsOnly = true
		page.Parse()
		// TODO: add page to categories
		// TODO: cache the error
		return DisplayError{Error: err.Error(), ParseError: true}
	}

	// if this is a draft and we're not serving drafts, pretend
	// that the page does not exist
	if !draftOK && page.Draft() {
		return DisplayError{Error: "Page has not yet been publised.", Draft: true}
	}

	// THIRD redirect check -
	// this is for pages we just parsed with @page.redirect
	if redir := page.Redirect(); redir != "" {
		return DisplayRedirect{Redirect: redir}
	}

	// generate HTML and metadata
	r.Generated = true
	r.Title = page.Title()
	r.FmtTitle = page.FmtTitle()
	r.Author = page.Author()
	r.Draft = page.Draft()
	r.ModUnix = page.Modified().Unix()
	r.Modified = httpdate.Time2Str(page.Modified())
	r.Content = page.HTML()
	r.CSS = page.CSS()

	// TODO: should we include the page object?
	// TODO: warnings

	// TODO: update categories and set to r.Categories

	// write cache file if enabled
	if dispErr := w.writePageCache(page, &r); dispErr != nil {
		return dispErr
	}

	// TODO: write search file if enabled

	return r
}

func (w *Wiki) writePageCache(page *wikifier.Page, r *DisplayPage) interface{} {

	// caching isn't enabled
	if !page.Opt.Page.EnableCache || page.CachePath() == "" {
		return nil
	}

	// open the cache file for writing
	cacheFile, err := os.Create(page.CachePath())
	defer cacheFile.Close()
	if err != nil {
		return DisplayError{
			Error:         "Could not write page cache file.",
			DetailedError: "Open '" + page.CachePath() + "' for write error: " + err.Error(),
		}
	}

	// generate page info
	info := pageJSONManifest{
		CSS:        page.CSS(),
		Categories: []string{}, // TODO
		Warnings:   []string{}, // TODO
		Error:      "",         // TODO
		PageInfo:   page.Info(),
	}

	// encode as json
	j, err := json.Marshal(info)
	if err != nil {
		return DisplayError{
			Error:         "Could not write page cache file.",
			DetailedError: "JSON encode error: " + err.Error(),
		}
	}

	// save prefixing data
	cacheFile.Write(j)
	cacheFile.Write([]byte{'\n'})

	// save content
	content := string(r.Content)
	cacheFile.WriteString(content)
	if len(content) != 0 && content[len(content)-1] != '\n' {
		cacheFile.Write([]byte{'\n'})
	}

	// update result with real cache modified times
	r.ModUnix = page.CacheModified().Unix()
	r.Modified = httpdate.Time2Str(page.CacheModified())
	r.CacheGenerated = true

	return nil // success
}

func (w *Wiki) displayCachedPage(page *wikifier.Page, r *DisplayPage, draftOK bool) interface{} {
	cacheModify := page.CacheModified()
	timeStr := httpdate.Time2Str(cacheModify)

	// the page's file is more recent than the cache file.
	// discard the outdated cached copy
	if page.Modified().After(cacheModify) {
		os.Remove(page.CachePath())
		return nil // OK
	}

	content := "<!-- cached page dated " + timeStr + " -->\n"

	// open cache file for reading
	cacheContent, err := ioutil.ReadFile(page.CachePath())
	if err != nil {
		return DisplayError{
			Error:         "Could not read page cache file.",
			DetailedError: "Open '" + page.CachePath() + "' for read error: " + err.Error(),
		}
	}

	// find the first line
	var jsonData []byte
	firstNL := bytes.IndexByte(cacheContent, '\n')
	if firstNL != -1 && firstNL < len(cacheContent) {
		jsonData = cacheContent[:firstNL+1]
		cacheContent = cacheContent[firstNL:]
	}

	// the rest is the html content
	content += string(cacheContent)

	// decode the manifest
	var info pageJSONManifest
	if err := json.Unmarshal(jsonData, &info); err != nil {
		return DisplayError{
			Error:         "Could not read page cache file.",
			DetailedError: "JSON encode error: " + err.Error(),
		}
	}

	// if this is a draft and we're not serving drafts, pretend
	// that the page does not exist
	if !draftOK && info.Draft {
		return DisplayError{Error: "Page has not yet been publised.", Draft: true}
	}

	// cached error
	if info.Error != "" {
		return DisplayError{Error: info.Error}
	}

	// cached redirect
	if info.Redirect != "" {
		return DisplayRedirect{Redirect: info.Redirect}
	}

	// update result with stuff from cache
	r.CreatedUnix = info.Created.Unix()
	r.Draft = info.Draft
	r.Author = info.Author
	r.Title = info.Title
	r.FmtTitle = info.FmtTitle
	r.FromCache = true
	r.Content = wikifier.HTML(content)
	r.ModUnix = cacheModify.Unix()
	r.Modified = httpdate.Time2Str(cacheModify)

	return nil // success
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"

	"github.com/pmylund/go-cache"
)

const (
	dataFileExt = ".data"
)

type Sample struct {
	ts int64
	v  float64
}

type perfDB struct {
	baseDir string
	cache   cache.Cache
	mu      sync.Mutex
}

func newPerfDB(baseDir string) (*perfDB, error) {
	if err := os.MkdirAll(baseDir, 0755); err != nil {
		logger.Critical("Failed to initialize datastore: %s", err)
		return nil, err
	}
	c := cache.New(cache.NoExpiration, cache.NoExpiration)
	return &perfDB{baseDir, *c, sync.Mutex{}}, nil
}

func (pdb *perfDB) getDirPath(dbname string) string {
	return filepath.Join(pdb.baseDir, dbname)
}

func (pdb *perfDB) getFilePath(dbname, metric string) string {
	dataDir := pdb.getDirPath(dbname)

	return filepath.Join(dataDir, metric+dataFileExt)
}

func (pdb *perfDB) isExist(dbname string) (bool, error) {
	dataDir := pdb.getDirPath(dbname)

	_, err := os.Stat(dataDir)
	if err == nil {
		return true, nil
	}
	return false, err
}

func (pdb *perfDB) listDatabases() ([]string, error) {
	files, err := ioutil.ReadDir(pdb.baseDir)
	if err != nil {
		return nil, err
	}
	databases := []string{}
	for _, f := range files {
		databases = append(databases, f.Name())
	}
	return databases, nil
}

func (pdb *perfDB) listMetrics(dbname string) ([]string, error) {
	dataDir := pdb.getDirPath(dbname)
	files, err := ioutil.ReadDir(dataDir)
	if err != nil {
		return nil, err
	}
	metrics := []string{}
	for _, f := range files {
		fname := strings.Replace(f.Name(), dataFileExt, "", 1)
		metrics = append(metrics, fname)
	}
	return metrics, nil
}

const tsOffset = 22

func (pdb *perfDB) addSample(dbname, metric string, sample Sample) error {
	dataDir := pdb.getDirPath(dbname)
	if err := os.MkdirAll(dataDir, 0775); err != nil {
		return err
	}

	dataFile := pdb.getFilePath(dbname, metric)

	file, err := os.OpenFile(dataFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		logger.Critical(err)
		return err
	}
	defer file.Close()

	pdb.mu.Lock()
	defer pdb.mu.Unlock()
	if _, err := fmt.Fprintf(file, "%22d %025.9f\n", sample.ts, sample.v); err != nil {
		logger.Critical(err)
		return err
	}

	return nil
}

const bufferSize = 1000

func fullScan(fileName string, done <-chan struct{}) (<-chan string, <-chan error) {
	samples := make(chan string, bufferSize)
	errc := make(chan error, 1)

	go func() {
		defer close(samples)
		defer close(errc)

		file, err := os.Open(fileName)
		if err != nil {
			errc <- err
			return
		}
		defer file.Close()

		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			select {
			case samples <- scanner.Text():
			case <-done:
				return
			}
		}

		if err := scanner.Err(); err != nil {
			errc <- err
		}
	}()
	return samples, errc
}

func parseSamples(records <-chan string) (<-chan Sample, <-chan error) {
	samples := make(chan Sample, bufferSize)
	errc := make(chan error, 1)

	go func() {
		defer close(samples)
		defer close(errc)

		var err error
		var value float64
		var ts int64
		for record := range records {
			if value, err = strconv.ParseFloat(record[tsOffset+1:], 64); err != nil {
				errc <- err
				break
			}
			if ts, err = strconv.ParseInt(strings.TrimSpace(record[:tsOffset]), 10, 64); err != nil {
				errc <- err
				break
			}
			samples <- Sample{ts, value}
		}
	}()
	return samples, errc
}

func mergeErrors(errcs ...<-chan error) error {
	for _, errc := range errcs {
		if err := <-errc; err != nil {
			return err
		}
	}
	return nil
}

func (pdb *perfDB) getRawValues(dbname, metric string) ([][]interface{}, error) {
	dataFile := pdb.getFilePath(dbname, metric)

	done := make(chan struct{}, 1)
	rawSamples, rawErrors := fullScan(dataFile, done)
	parsedSamples, parsedErrors := parseSamples(rawSamples)

	values := [][]interface{}{}
	for sample := range parsedSamples {
		values = append(values, []interface{}{sample.ts, sample.v})
	}

	done <- struct{}{}
	if err := mergeErrors(rawErrors, parsedErrors); err != nil {
		return nil, err
	}
	return values, nil
}

func metricHash(filePath string) (string, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return "", err
	}
	defer file.Close()

	info, err := file.Stat()
	if err != nil {
		return "", err
	}

	return fmt.Sprintf("%s%d", filePath, info.Size()), nil
}

func (pdb *perfDB) getSummary(dbname, metric string) (map[string]interface{}, error) {
	var summary map[string]interface{}

	dataFile := pdb.getFilePath(dbname, metric)

	hash, err := metricHash(dataFile)
	if err != nil {
		return nil, err
	}

	if cachedData, found := pdb.cache.Get(string(hash)); found {
		return cachedData.(map[string]interface{}), nil
	}

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

	rawSamples, rawErrors := fullScan(dataFile, done)
	parsedSamples, parsedErrors := parseSamples(rawSamples)

	values := []float64{}
	sum := 0.0
	for sample := range parsedSamples {
		sum += sample.v
		values = append(values, sample.v)
	}

	done <- struct{}{}
	if err := mergeErrors(rawErrors, parsedErrors); err != nil {
		return nil, err
	}

	count := len(values)
	sort.Float64s(values)

	summary = map[string]interface{}{
		"max":   values[count-1],
		"min":   values[0],
		"count": count,
		"avg":   sum / float64(count),
	}

	for _, percentile := range []float64{0.5, 0.8, 0.9, 0.95, 0.99, 0.999} {
		var pIdx int
		if count > 1 {
			pIdx = int(float64(count)*percentile) - 1
		}
		p := fmt.Sprintf("p%v", percentile*100)
		summary[p] = values[pIdx]
	}

	pdb.cache.Set(hash, summary, cache.NoExpiration)
	return summary, nil
}

type parsedSample struct {
	ts int64
	v  float64
}

func parseSamplesWithTimestamp(records <-chan string) (<-chan parsedSample, <-chan error) {
	samples := make(chan parsedSample, bufferSize)
	errc := make(chan error, 1)

	go func() {
		defer close(samples)
		defer close(errc)

		var err error
		var ts int64
		var value float64
		for record := range records {
			if value, err = strconv.ParseFloat(record[tsOffset+1:], 64); err != nil {
				errc <- err
				break
			}
			if ts, err = strconv.ParseInt(strings.TrimSpace(record[:tsOffset]), 10, 64); err != nil {
				errc <- err
				break
			}
			samples <- parsedSample{ts, value}
		}
	}()
	return samples, errc
}

func (pdb *perfDB) getHeatMap(dbname, metric string) (*heatMap, error) {
	dataFile := pdb.getFilePath(dbname, metric)

	hm := newHeatMap()
	hm.MinTS = int64(^uint64(0) >> 1)

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

	rawSamples, rawErrors := fullScan(dataFile, done)
	parsedSamples, parsedErrors := parseSamplesWithTimestamp(rawSamples)

	samples := []parsedSample{}
	for sample := range parsedSamples {
		hm.MaxValue = math.Max(hm.MaxValue, sample.v)
		if sample.ts < hm.MinTS {
			hm.MinTS = sample.ts
		} else if sample.ts > hm.MaxTS {
			hm.MaxTS = sample.ts
		}
		samples = append(samples, sample)
	}

	done <- struct{}{}
	if err := mergeErrors(rawErrors, parsedErrors); err != nil {
		return nil, err
	}

	for _, sample := range samples {
		x := math.Floor(heatMapWidth * float64(sample.ts-hm.MinTS) / float64(hm.MaxTS-hm.MinTS))
		y := math.Floor(heatMapHeight * sample.v / hm.MaxValue)
		if x == heatMapWidth {
			x--
		}
		if y == heatMapHeight {
			y--
		}
		hm.Map[int(y)][int(x)]++
		if hm.Map[int(y)][int(x)] > hm.maxDensity {
			hm.maxDensity = hm.Map[int(y)][int(x)]
		}
	}
	return hm, nil
}
<commit_msg>Do not open file for hashing<commit_after>package main

import (
	"bufio"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"

	"github.com/pmylund/go-cache"
)

const (
	dataFileExt = ".data"
)

type Sample struct {
	ts int64
	v  float64
}

type perfDB struct {
	baseDir string
	cache   cache.Cache
	mu      sync.Mutex
}

func newPerfDB(baseDir string) (*perfDB, error) {
	if err := os.MkdirAll(baseDir, 0755); err != nil {
		logger.Critical("Failed to initialize datastore: %s", err)
		return nil, err
	}
	c := cache.New(cache.NoExpiration, cache.NoExpiration)
	return &perfDB{baseDir, *c, sync.Mutex{}}, nil
}

func (pdb *perfDB) getDirPath(dbname string) string {
	return filepath.Join(pdb.baseDir, dbname)
}

func (pdb *perfDB) getFilePath(dbname, metric string) string {
	dataDir := pdb.getDirPath(dbname)

	return filepath.Join(dataDir, metric+dataFileExt)
}

func (pdb *perfDB) isExist(dbname string) (bool, error) {
	dataDir := pdb.getDirPath(dbname)

	_, err := os.Stat(dataDir)
	if err == nil {
		return true, nil
	}
	return false, err
}

func (pdb *perfDB) listDatabases() ([]string, error) {
	files, err := ioutil.ReadDir(pdb.baseDir)
	if err != nil {
		return nil, err
	}
	databases := []string{}
	for _, f := range files {
		databases = append(databases, f.Name())
	}
	return databases, nil
}

func (pdb *perfDB) listMetrics(dbname string) ([]string, error) {
	dataDir := pdb.getDirPath(dbname)
	files, err := ioutil.ReadDir(dataDir)
	if err != nil {
		return nil, err
	}
	metrics := []string{}
	for _, f := range files {
		fname := strings.Replace(f.Name(), dataFileExt, "", 1)
		metrics = append(metrics, fname)
	}
	return metrics, nil
}

const tsOffset = 22

func (pdb *perfDB) addSample(dbname, metric string, sample Sample) error {
	dataDir := pdb.getDirPath(dbname)
	if err := os.MkdirAll(dataDir, 0775); err != nil {
		return err
	}

	dataFile := pdb.getFilePath(dbname, metric)

	file, err := os.OpenFile(dataFile, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0644)
	if err != nil {
		logger.Critical(err)
		return err
	}
	defer file.Close()

	pdb.mu.Lock()
	defer pdb.mu.Unlock()
	if _, err := fmt.Fprintf(file, "%22d %025.9f\n", sample.ts, sample.v); err != nil {
		logger.Critical(err)
		return err
	}

	return nil
}

const bufferSize = 1000

func fullScan(fileName string, done <-chan struct{}) (<-chan string, <-chan error) {
	samples := make(chan string, bufferSize)
	errc := make(chan error, 1)

	go func() {
		defer close(samples)
		defer close(errc)

		file, err := os.Open(fileName)
		if err != nil {
			errc <- err
			return
		}
		defer file.Close()

		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			select {
			case samples <- scanner.Text():
			case <-done:
				return
			}
		}

		if err := scanner.Err(); err != nil {
			errc <- err
		}
	}()
	return samples, errc
}

func parseSamples(records <-chan string) (<-chan Sample, <-chan error) {
	samples := make(chan Sample, bufferSize)
	errc := make(chan error, 1)

	go func() {
		defer close(samples)
		defer close(errc)

		var err error
		var value float64
		var ts int64
		for record := range records {
			if value, err = strconv.ParseFloat(record[tsOffset+1:], 64); err != nil {
				errc <- err
				break
			}
			if ts, err = strconv.ParseInt(strings.TrimSpace(record[:tsOffset]), 10, 64); err != nil {
				errc <- err
				break
			}
			samples <- Sample{ts, value}
		}
	}()
	return samples, errc
}

func mergeErrors(errcs ...<-chan error) error {
	for _, errc := range errcs {
		if err := <-errc; err != nil {
			return err
		}
	}
	return nil
}

func (pdb *perfDB) getRawValues(dbname, metric string) ([][]interface{}, error) {
	dataFile := pdb.getFilePath(dbname, metric)

	done := make(chan struct{}, 1)
	rawSamples, rawErrors := fullScan(dataFile, done)
	parsedSamples, parsedErrors := parseSamples(rawSamples)

	values := [][]interface{}{}
	for sample := range parsedSamples {
		values = append(values, []interface{}{sample.ts, sample.v})
	}

	done <- struct{}{}
	if err := mergeErrors(rawErrors, parsedErrors); err != nil {
		return nil, err
	}
	return values, nil
}

func metricHash(filePath string) (string, error) {
	info, err := os.Stat(filePath)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("%s%d", filePath, info.Size()), nil
}

func (pdb *perfDB) getSummary(dbname, metric string) (map[string]interface{}, error) {
	var summary map[string]interface{}

	dataFile := pdb.getFilePath(dbname, metric)

	hash, err := metricHash(dataFile)
	if err != nil {
		return nil, err
	}

	if cachedData, found := pdb.cache.Get(string(hash)); found {
		return cachedData.(map[string]interface{}), nil
	}

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

	rawSamples, rawErrors := fullScan(dataFile, done)
	parsedSamples, parsedErrors := parseSamples(rawSamples)

	values := []float64{}
	sum := 0.0
	for sample := range parsedSamples {
		sum += sample.v
		values = append(values, sample.v)
	}

	done <- struct{}{}
	if err := mergeErrors(rawErrors, parsedErrors); err != nil {
		return nil, err
	}

	count := len(values)
	sort.Float64s(values)

	summary = map[string]interface{}{
		"max":   values[count-1],
		"min":   values[0],
		"count": count,
		"avg":   sum / float64(count),
	}

	for _, percentile := range []float64{0.5, 0.8, 0.9, 0.95, 0.99, 0.999} {
		var pIdx int
		if count > 1 {
			pIdx = int(float64(count)*percentile) - 1
		}
		p := fmt.Sprintf("p%v", percentile*100)
		summary[p] = values[pIdx]
	}

	pdb.cache.Set(hash, summary, cache.NoExpiration)
	return summary, nil
}

type parsedSample struct {
	ts int64
	v  float64
}

func parseSamplesWithTimestamp(records <-chan string) (<-chan parsedSample, <-chan error) {
	samples := make(chan parsedSample, bufferSize)
	errc := make(chan error, 1)

	go func() {
		defer close(samples)
		defer close(errc)

		var err error
		var ts int64
		var value float64
		for record := range records {
			if value, err = strconv.ParseFloat(record[tsOffset+1:], 64); err != nil {
				errc <- err
				break
			}
			if ts, err = strconv.ParseInt(strings.TrimSpace(record[:tsOffset]), 10, 64); err != nil {
				errc <- err
				break
			}
			samples <- parsedSample{ts, value}
		}
	}()
	return samples, errc
}

func (pdb *perfDB) getHeatMap(dbname, metric string) (*heatMap, error) {
	dataFile := pdb.getFilePath(dbname, metric)

	hm := newHeatMap()
	hm.MinTS = int64(^uint64(0) >> 1)

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

	rawSamples, rawErrors := fullScan(dataFile, done)
	parsedSamples, parsedErrors := parseSamplesWithTimestamp(rawSamples)

	samples := []parsedSample{}
	for sample := range parsedSamples {
		hm.MaxValue = math.Max(hm.MaxValue, sample.v)
		if sample.ts < hm.MinTS {
			hm.MinTS = sample.ts
		} else if sample.ts > hm.MaxTS {
			hm.MaxTS = sample.ts
		}
		samples = append(samples, sample)
	}

	done <- struct{}{}
	if err := mergeErrors(rawErrors, parsedErrors); err != nil {
		return nil, err
	}

	for _, sample := range samples {
		x := math.Floor(heatMapWidth * float64(sample.ts-hm.MinTS) / float64(hm.MaxTS-hm.MinTS))
		y := math.Floor(heatMapHeight * sample.v / hm.MaxValue)
		if x == heatMapWidth {
			x--
		}
		if y == heatMapHeight {
			y--
		}
		hm.Map[int(y)][int(x)]++
		if hm.Map[int(y)][int(x)] > hm.maxDensity {
			hm.maxDensity = hm.Map[int(y)][int(x)]
		}
	}
	return hm, nil
}
<|endoftext|>