<commit_before>// +build examples

package main

import (
	"encoding/hex"
	"fmt"
	"gopkg.in/jcmturner/gokrb5.v5/client"
	"gopkg.in/jcmturner/gokrb5.v5/config"
	"gopkg.in/jcmturner/gokrb5.v5/credentials"
	"gopkg.in/jcmturner/gokrb5.v5/keytab"
	"gopkg.in/jcmturner/gokrb5.v5/service"
	"gopkg.in/jcmturner/gokrb5.v5/testdata"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
)

func main() {
	s := httpServer()
	defer s.Close()

	b, _ := hex.DecodeString(testdata.TESTUSER1_KEYTAB)
	kt, _ := keytab.Parse(b)
	c, _ := config.NewConfigFromString(testdata.TEST_KRB5CONF)
	cl := client.NewClientWithKeytab("testuser1", "TEST.GOKRB5", kt)
	cl.WithConfig(c)
	httpRequest(s.URL, cl)

	b, _ = hex.DecodeString(testdata.TESTUSER2_KEYTAB)
	kt, _ = keytab.Parse(b)
	c, _ = config.NewConfigFromString(testdata.TEST_KRB5CONF)
	cl = client.NewClientWithKeytab("testuser2", "TEST.GOKRB5", kt)
	cl.WithConfig(c)
	httpRequest(s.URL, cl)

	//httpRequest("http://host.test.gokrb5/index.html")
}

func httpRequest(url string, cl client.Client) {
	l := log.New(os.Stderr, "GOKRB5 Client: ", log.Ldate|log.Ltime|log.Lshortfile)

	err := cl.Login()
	if err != nil {
		l.Printf("Error on AS_REQ: %v\n", err)
	}
	r, _ := http.NewRequest("GET", url, nil)
	err = cl.SetSPNEGOHeader(r, "HTTP/host.test.gokrb5")
	if err != nil {
		l.Printf("Error setting client SPNEGO header: %v", err)
	}
	httpResp, err := http.DefaultClient.Do(r)
	if err != nil {
		l.Printf("Request error: %v\n", err)
	}
	fmt.Fprintf(os.Stdout, "Response Code: %v\n", httpResp.StatusCode)
	content, _ := ioutil.ReadAll(httpResp.Body)
	fmt.Fprintf(os.Stdout, "Response Body:\n%s\n", content)
}

func httpServer() *httptest.Server {
	l := log.New(os.Stderr, "GOKRB5 Service: ", log.Ldate|log.Ltime|log.Lshortfile)
	b, _ := hex.DecodeString(testdata.HTTP_KEYTAB)
	kt, _ := keytab.Parse(b)
	th := http.HandlerFunc(testAppHandler)
	s := httptest.NewServer(service.SPNEGOKRB5Authenticate(th, kt, "", false, l))
	return s
}

func testAppHandler(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	fmt.Fprint(w, "\n

TEST.GOKRB5 Handler

\n")
	if validuser, ok := ctx.Value(service.CTXKeyAuthenticated).(bool); ok && validuser {
		if creds, ok := ctx.Value(service.CTXKeyCredentials).(credentials.Credentials); ok {
			fmt.Fprintf(w, "

Authenticed user: %s
\n", creds.Username)
			fmt.Fprintf(w, "
User's realm: %s

\n", creds.Realm)
		}

	} else {
		w.WriteHeader(http.StatusUnauthorized)
		fmt.Fprint(w, "Authentication failed")
	}
	fmt.Fprint(w, "")
	return
}
<commit_msg>example<commit_after>// +build examples

package main

import (
	"encoding/hex"
	"fmt"
	"gopkg.in/jcmturner/gokrb5.v5/client"
	"gopkg.in/jcmturner/gokrb5.v5/config"
	"gopkg.in/jcmturner/gokrb5.v5/credentials"
	"gopkg.in/jcmturner/gokrb5.v5/keytab"
	"gopkg.in/jcmturner/gokrb5.v5/service"
	"gopkg.in/jcmturner/gokrb5.v5/testdata"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
)

func main() {
	s := httpServer()
	defer s.Close()

	b, _ := hex.DecodeString(testdata.TESTUSER1_KEYTAB)
	kt, _ := keytab.Parse(b)
	c, _ := config.NewConfigFromString(testdata.TEST_KRB5CONF)
	c.LibDefaults.NoAddresses = true
	cl := client.NewClientWithKeytab("testuser1", "TEST.GOKRB5", kt)
	cl.WithConfig(c)
	httpRequest(s.URL, cl)

	b, _ = hex.DecodeString(testdata.TESTUSER2_KEYTAB)
	kt, _ = keytab.Parse(b)
	c, _ = config.NewConfigFromString(testdata.TEST_KRB5CONF)
	c.LibDefaults.NoAddresses = true
	cl = client.NewClientWithKeytab("testuser2", "TEST.GOKRB5", kt)
	cl.WithConfig(c)
	httpRequest(s.URL, cl)
}

func httpRequest(url string, cl client.Client) {
	l := log.New(os.Stderr, "GOKRB5 Client: ", log.Ldate|log.Ltime|log.Lshortfile)

	err := cl.Login()
	if err != nil {
		l.Printf("Error on AS_REQ: %v\n", err)
	}
	r, _ := http.NewRequest("GET", url, nil)
	err = cl.SetSPNEGOHeader(r, "HTTP/host.test.gokrb5")
	if err != nil {
		l.Printf("Error setting client SPNEGO header: %v", err)
	}
	httpResp, err := http.DefaultClient.Do(r)
	if err != nil {
		l.Printf("Request error: %v\n", err)
	}
	fmt.Fprintf(os.Stdout, "Response Code: %v\n", httpResp.StatusCode)
	content, _ := ioutil.ReadAll(httpResp.Body)
	fmt.Fprintf(os.Stdout, "Response Body:\n%s\n", content)
}

func httpServer() *httptest.Server {
	l := log.New(os.Stderr, "GOKRB5 Service: ", log.Ldate|log.Ltime|log.Lshortfile)
	b, _ := hex.DecodeString(testdata.HTTP_KEYTAB)
	kt, _ := keytab.Parse(b)
	th := http.HandlerFunc(testAppHandler)
	s := httptest.NewServer(service.SPNEGOKRB5Authenticate(th, kt, "", false, l))
	return s
}

func testAppHandler(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	fmt.Fprint(w, "\n

TEST.GOKRB5 Handler

\n")
	if validuser, ok := ctx.Value(service.CTXKeyAuthenticated).(bool); ok && validuser {
		if creds, ok := ctx.Value(service.CTXKeyCredentials).(credentials.Credentials); ok {
			fmt.Fprintf(w, "

Authenticed user: %s
\n", creds.Username)
			fmt.Fprintf(w, "
User's realm: %s

\n", creds.Realm)
		}

	} else {
		w.WriteHeader(http.StatusUnauthorized)
		fmt.Fprint(w, "Authentication failed")
	}
	fmt.Fprint(w, "")
	return
}
<|endoftext|>


<commit_before>package s3test

import (

	"github.com/stretchr/testify/suite"
	"fmt"

	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	//"github.com/aws/aws-sdk-go/awstesting/unit"
	"github.com/spf13/viper"
)

type HostStyleSuite struct {
    suite.Suite
}

type s3BucketTest struct {
	bucket  string
	url     string
	errCode string
}

var domain = fmt.Sprintf( "http://%s", viper.GetString("s3main.endpoint") )

func getUrl(bktname string)(string){
	url := fmt.Sprintf("%s%s", domain, bktname)

	return url
}

// func (suite *HostStyleSuite) TestAccelerateNoSSLBucketBuild() {

// 	assert := suite

// 	tests := []s3BucketTest{
// 		{"a.b.c", getUrl("a.b.c"), ""},
// 		{"a..bc", getUrl("a..bc"), ""},
// 	}

// 	for _, test := range tests {
// 		req, _ := svc.ListObjectsRequest(&s3.ListObjectsInput{Bucket: &test.bucket})
// 		req.Build()
// 		assert.Equal (test.url, req.HTTPRequest.URL.String())
// 	}
// }

func (suite *HostStyleSuite) TestHostStyleBucketBuildNoSSL() {

	assert := suite

	tests := []s3BucketTest{
		{"abc", getUrl("abc/key"), ""},
		// {"a.b.c", getUrl("a.b.c"), ""},
		// {"a$b$c", getUrl("%7BBucket%7D"), "InvalidParameterException"},
	}

	for _, test := range tests {

		input := &s3.CopyObjectInput{
			Bucket:         aws.String("bucket"),
			CopySource:     aws.String("bucket/source"),
			Key:            aws.String("key"),
		}

		input.SetWebsiteRedirectLocation(string(test.url))
		req, _ := svc.CopyObjectRequest(input)
		err := req.Send()

		assert.Equal(test.url, req.HTTPRequest.URL.String())

		if err != nil {
			if awsErr, ok := err.(awserr.Error); ok {

				assert.Equal(awsErr.Code(), test.errCode)
				assert.Equal(awsErr.Message(), "")
			}
		}
	}
}

<commit_msg>host style<commit_after>package s3test

import (

	"github.com/stretchr/testify/suite"
	"fmt"

	"github.com/aws/aws-sdk-go/service/s3"
	//"github.com/aws/aws-sdk-go/aws"
	//"github.com/aws/aws-sdk-go/awstesting/unit"
	"github.com/spf13/viper"
)

type HostStyleSuite struct {
    suite.Suite
}

type s3BucketTest struct {
	bucket  string
	url     string
	errCode string
}

var domain = fmt.Sprintf( "http://%s", viper.GetString("s3main.endpoint") )

func getUrl(bktname string)(string){
	url := fmt.Sprintf("%s%s", domain, bktname)

	return url
}

func (suite *HostStyleSuite) TestAccelerateNoSSLBucketBuild() {

	assert := suite

	tests := []s3BucketTest{
		{"a.b.c", getUrl("a.b.c"), ""},
		{"a..bc", getUrl("a..bc"), ""},
	}

	for _, test := range tests {
		req, _ := svc.ListObjectsRequest(&s3.ListObjectsInput{Bucket: &test.bucket})
		req.Build()
		assert.Equal (test.url, req.HTTPRequest.URL.String())
	}
}

func (suite *HostStyleSuite) TestHostStyleBucketBuildNoSSL() {

	//svc := s3.New(unit.Session, &aws.Config{DisableSSL: aws.Bool(true)})

	assert := suite

	tests := []s3BucketTest{
		{"abc", getUrl("abc"), ""},
		{"a.b.c", getUrl("a.b.c"), ""},
		{"a$b$c", getUrl("a%24b%24c"), "InvalidParameterException"},
	}

	for _, test := range tests {
		req, _ := svc.ListObjectsRequest(&s3.ListObjectsInput{Bucket: &test.bucket})
		req.Build()
		assert.Equal (test.url, req.HTTPRequest.URL.String())
	}
}

<|endoftext|>


<commit_before>// file test must be end by _test example "koans_test.go"
// To run test suite execute "go test" command

package koans_test

import(
  "fmt" //
  "testing" // testing package : https://golang.org/pkg/testing/
)

// All test must be start with TestXxx(*testing.T)
func TestSample(t *testing.T) {
  // To signal failure use FailNow, Error, Fail, etc. func

  t.Skip("skipping test message") // To skip a test
}

// Benchmarks
// The benchmarks func must use this pattern
// func BenchmarkXxx(*testing.B)
func BenchmarkHello(b *testing.B) {
  for i := 0; i < b.N; i++ {
    fmt.Println("hello world")
  }
}

// Pointer

// Notes
// Go have pointer but go has no pointer arithmetic :(

func TestPointer(t *testing.T) {
  // Use {} to isolate context into go func
  // Relly usefull to have many idependant
  // assertation into func Test
  {
    x := 3
    y := x // b is a copy of a
    // All assignments in golang are copy operation

    y++

    if x == y {
      t.FailNow()
    }
  }
  {
    x := 3
    y := &x // use & to access to adress of a pointer

    if *y != 3 {
      t.FailNow()
    }

    // y is a pointer of x
    // Use * symbol to de-referencing
    *y = *y + 2

    if x != 5 {
      t.FailNow()
    }
  }
  {
    incr := func(i int) int {
      i = i + 1
      return i
    }

    x := 1
    incr(x)

    if x != 1 {
      t.FailNow()
    }
  }
  {
    incr := func(i *int) *int {
      *i++
      return i
    }

    x := 1
    incr(&x)
    if x != 2 {
      t.FailNow()
    }
  }
  {
    x := 1
    p := &x
    pp := &p
    cp := **pp

    if **pp != x {
      t.FailNow()
    }

    **pp = 4
    if x != 4 {
      t.FailNow()
    }

    if cp != 1 {
      t.FailNow()
    }
  }

}
<commit_msg>Add assert helper<commit_after>// file test must be end by _test example "koans_test.go"
// To run test suite execute "go test" command

package koans_test

import(
  "fmt" //
  "testing" // testing package : https://golang.org/pkg/testing/
)

// All test must be start with TestXxx(*testing.T)
func TestSample(t *testing.T) {
  // To signal failure use FailNow, Error, Fail, etc. func

  t.Skip("skipping test message") // To skip a test
}

// Benchmarks
// The benchmarks func must use this pattern
// func BenchmarkXxx(*testing.B)
func BenchmarkHello(b *testing.B) {
  for i := 0; i < b.N; i++ {
    fmt.Println("hello world")
  }
}

// Little helper to write test
func Assert(t *testing.T, assertation bool) {
  if assertation {
    fmt.Println("*")
  } else {
    t.FailNow()
  }
}

// Pointer
// Notes
// Go have pointer but go has no pointer arithmetic :(

func TestPointer(t *testing.T) {
  // Use {} to isolate context into go func
  // Relly usefull to have many idependant
  // assertation into func Test
  {
    x := 3
    y := x // b is a copy of a
    // All assignments in golang are copy operation

    y++

    Assert(t, x != y)
  }
  {
    x := 3
    y := &x // use & to access to adress of a pointer

    Assert(t, *y == 3)

    // y is a pointer of x
    // Use * symbol to de-referencing
    *y = *y + 2

    Assert(t, x == 5)
  }
  {
    incr := func(i int) int {
      i = i + 1
      return i
    }

    x := 1
    incr(x)

    Assert(t, x == 1)
  }
  {
    incr := func(i *int) *int {
      *i++
      return i
    }

    x := 1
    incr(&x)

    Assert(t, x == 2)
  }
  {
    x := 1
    p := &x
    pp := &p
    cp := **pp

    Assert(t, **pp == x)

    // Modify value of x
    **pp = 4
    Assert(t, x == 4)
    Assert(t, cp == 1)
  }
}
<|endoftext|>


<commit_before>// (c) 2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package socket

import (
	"encoding/binary"
	"errors"
	"io"
	"net"
	"os"
	"sync"
	"syscall"

	"github.com/ava-labs/gecko/utils/logging"
	"github.com/ava-labs/gecko/utils/wrappers"
)

var (
	// ErrTimeout is returned when a socket operation times out
	ErrTimeout = errors.New("timeout")
)

// Socket manages sending messages over a socket to many subscribed clients
type Socket struct {
	log logging.Logger

	addr     string
	accept   acceptFn
	connLock *sync.RWMutex
	conns    []net.Conn

	quitCh chan struct{}
	doneCh chan struct{}
}

// NewSocket creates a new socket object for the given address. It does not open
// the socket until Listen is called.
func NewSocket(addr string, log logging.Logger) *Socket {
	return &Socket{
		log: log,

		addr:     addr,
		accept:   accept,
		connLock: &sync.RWMutex{},

		quitCh: make(chan struct{}),
		doneCh: make(chan struct{}),
	}
}

// Listen starts listening on the socket for new connection
func (s *Socket) Listen() error {
	l, err := listen(s.addr)
	if err != nil {
		return err
	}

	// Start a loop that accepts new connections to told to quit
	go func() {
		for {
			select {
			case <-s.quitCh:
				close(s.doneCh)
				return
			default:
				s.accept(s, l)
			}
		}
	}()

	return nil
}

// Send writes the given message to all connection clients
func (s *Socket) Send(msg []byte) error {
	// Prefix the message with an 8 byte length
	lenBytes := make([]byte, 8)
	binary.BigEndian.PutUint64(lenBytes, uint64(len(msg)))
	msg = append(lenBytes, msg...)

	// Get a copy of connections
	s.connLock.RLock()
	conns := make([]net.Conn, len(s.conns))
	copy(conns, s.conns)
	s.connLock.RUnlock()

	// Write to each connection
	var err error
	for _, conn := range conns {
		if _, err = conn.Write(msg); err != nil {
			return err
		}
	}
	return nil
}

// Close closes the socket by cutting off new connections, closing all
// existing ones, and then zero'ing out the connection pool
func (s *Socket) Close() error {
	// Signal to the event loop to stop and wait for it to signal back
	close(s.quitCh)
	<-s.doneCh

	// Zero out the connection pool but save a reference so we can close them all
	s.connLock.Lock()
	conns := s.conns
	s.conns = nil
	s.connLock.Unlock()

	// Close all connections that were open at the time of shutdown
	errs := wrappers.Errs{}
	for _, conn := range conns {
		errs.Add(conn.Close())
	}
	return errs.Err
}

// Client is a read-only connection to a socket
type Client struct {
	net.Conn
}

// Recv waits for a message from the socket. It's guaranteed to either return a
// complete message or an error
func (c *Client) Recv() ([]byte, error) {
	// Read length
	var sz int64
	if err := binary.Read(c.Conn, binary.BigEndian, &sz); err != nil {
		if isTimeoutError(err) {
			return nil, ErrTimeout
		}
		return nil, err
	}

	// Create buffer for entire message and read it all in
	msg := make([]byte, sz)
	if _, err := io.ReadFull(c.Conn, msg); err != nil {
		if isTimeoutError(err) {
			return nil, ErrTimeout
		}
		return nil, err
	}

	return msg, nil
}

// Close closes the underlying socket connection
func (c *Client) Close() error {
	return c.Conn.Close()
}

// acceptFn takes accepts connections from a Listener and gives them to a Socket
type acceptFn func(*Socket, net.Listener)

// accept is the default acceptFn for sockets. It accepts the next connection
// from the given listen and adds it to the Socket's connection list
func accept(s *Socket, l net.Listener) {
	conn, err := l.Accept()
	if err != nil {
		s.log.Error("socket accept error: %s", err.Error())
	}
	s.connLock.Lock()
	s.conns = append(s.conns, conn)
	s.connLock.Unlock()
}

// isTimeoutError checks if an error is a timeout as per the net.Error interface
func isTimeoutError(err error) bool {
	iErr, ok := err.(net.Error)
	if !ok {
		return false
	}
	return iErr.Timeout()
}

// isSyscallError checks if an error is one of the given syscall.Errno codes
func isSyscallError(err error, codes ...syscall.Errno) bool {
	opErr, ok := err.(*net.OpError)
	if !ok {
		return false
	}
	syscallErr, ok := opErr.Err.(*os.SyscallError)
	if !ok {
		return false
	}
	errno, ok := syscallErr.Err.(syscall.Errno)
	if !ok {
		return false
	}
	for _, code := range codes {
		if errno == code {
			return true
		}
	}
	return false
}
<commit_msg>CLEANUP: Fix comment.<commit_after>// (c) 2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package socket

import (
	"encoding/binary"
	"errors"
	"io"
	"net"
	"os"
	"sync"
	"syscall"

	"github.com/ava-labs/gecko/utils/logging"
	"github.com/ava-labs/gecko/utils/wrappers"
)

var (
	// ErrTimeout is returned when a socket operation times out
	ErrTimeout = errors.New("timeout")
)

// Socket manages sending messages over a socket to many subscribed clients
type Socket struct {
	log logging.Logger

	addr     string
	accept   acceptFn
	connLock *sync.RWMutex
	conns    []net.Conn

	quitCh chan struct{}
	doneCh chan struct{}
}

// NewSocket creates a new socket object for the given address. It does not open
// the socket until Listen is called.
func NewSocket(addr string, log logging.Logger) *Socket {
	return &Socket{
		log: log,

		addr:     addr,
		accept:   accept,
		connLock: &sync.RWMutex{},

		quitCh: make(chan struct{}),
		doneCh: make(chan struct{}),
	}
}

// Listen starts listening on the socket for new connection
func (s *Socket) Listen() error {
	l, err := listen(s.addr)
	if err != nil {
		return err
	}

	// Start a loop that accepts new connections until told to quit
	go func() {
		for {
			select {
			case <-s.quitCh:
				close(s.doneCh)
				return
			default:
				s.accept(s, l)
			}
		}
	}()

	return nil
}

// Send writes the given message to all connection clients
func (s *Socket) Send(msg []byte) error {
	// Prefix the message with an 8 byte length
	lenBytes := make([]byte, 8)
	binary.BigEndian.PutUint64(lenBytes, uint64(len(msg)))
	msg = append(lenBytes, msg...)

	// Get a copy of connections
	s.connLock.RLock()
	conns := make([]net.Conn, len(s.conns))
	copy(conns, s.conns)
	s.connLock.RUnlock()

	// Write to each connection
	var err error
	for _, conn := range conns {
		if _, err = conn.Write(msg); err != nil {
			return err
		}
	}
	return nil
}

// Close closes the socket by cutting off new connections, closing all
// existing ones, and then zero'ing out the connection pool
func (s *Socket) Close() error {
	// Signal to the event loop to stop and wait for it to signal back
	close(s.quitCh)
	<-s.doneCh

	// Zero out the connection pool but save a reference so we can close them all
	s.connLock.Lock()
	conns := s.conns
	s.conns = nil
	s.connLock.Unlock()

	// Close all connections that were open at the time of shutdown
	errs := wrappers.Errs{}
	for _, conn := range conns {
		errs.Add(conn.Close())
	}
	return errs.Err
}

// Client is a read-only connection to a socket
type Client struct {
	net.Conn
}

// Recv waits for a message from the socket. It's guaranteed to either return a
// complete message or an error
func (c *Client) Recv() ([]byte, error) {
	// Read length
	var sz int64
	if err := binary.Read(c.Conn, binary.BigEndian, &sz); err != nil {
		if isTimeoutError(err) {
			return nil, ErrTimeout
		}
		return nil, err
	}

	// Create buffer for entire message and read it all in
	msg := make([]byte, sz)
	if _, err := io.ReadFull(c.Conn, msg); err != nil {
		if isTimeoutError(err) {
			return nil, ErrTimeout
		}
		return nil, err
	}

	return msg, nil
}

// Close closes the underlying socket connection
func (c *Client) Close() error {
	return c.Conn.Close()
}

// acceptFn takes accepts connections from a Listener and gives them to a Socket
type acceptFn func(*Socket, net.Listener)

// accept is the default acceptFn for sockets. It accepts the next connection
// from the given listen and adds it to the Socket's connection list
func accept(s *Socket, l net.Listener) {
	conn, err := l.Accept()
	if err != nil {
		s.log.Error("socket accept error: %s", err.Error())
	}
	s.connLock.Lock()
	s.conns = append(s.conns, conn)
	s.connLock.Unlock()
}

// isTimeoutError checks if an error is a timeout as per the net.Error interface
func isTimeoutError(err error) bool {
	iErr, ok := err.(net.Error)
	if !ok {
		return false
	}
	return iErr.Timeout()
}

// isSyscallError checks if an error is one of the given syscall.Errno codes
func isSyscallError(err error, codes ...syscall.Errno) bool {
	opErr, ok := err.(*net.OpError)
	if !ok {
		return false
	}
	syscallErr, ok := opErr.Err.(*os.SyscallError)
	if !ok {
		return false
	}
	errno, ok := syscallErr.Err.(syscall.Errno)
	if !ok {
		return false
	}
	for _, code := range codes {
		if errno == code {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package magicsql

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
)

// Querier defines an interface for top-level sql types that can run SQL and
// prepare statements
type Querier interface {
	Query(string, ...interface{}) (*sql.Rows, error)
	Exec(string, ...interface{}) (sql.Result, error)
	Prepare(string) (*sql.Stmt, error)
}

// Operation represents a short-lived single-purpose combination of database
// calls.  On the first failure, its internal error is set, which all
// "children" (statements, transactions, etc) will see.  All children will
// refuse to perform any functions once an error has occurred, making it safe
// to perform a chain of related database calls and only check for an error
// when it makes sense.
//
// When a transaction is started, the operation will route all database calls
// through the transaction instead of the global database handler.  At this
// time, only one transaction at a time is supported (i.e., no nesting
// transactions).
type Operation struct {
	parent *DB
	err    error
	tx     *sql.Tx
	q      Querier
	Dbg    bool
}

// NewOperation creates an operation in its default state: its parent is the
// passed-in DB instance, and it defaults to using direct database calls until
// a transaction is started.
func NewOperation(db *DB) *Operation {
	var o = &Operation{parent: db}
	o.q = o.parent.db

	return o
}

// Err returns the *first* error which occurred on any database call owned by the Operation
func (op *Operation) Err() error {
	return op.err
}

// SetErr tells the Operation to stop handling any more queries.  It shouldn't
// usually be called directly, but it can be if you need to tell the object
// "here's a thing that may be an error; don't do any more work if it is".
func (op *Operation) SetErr(err error) {
	if op.Err() != nil {
		return
	}

	op.err = err
}

// Query wraps sql's Query, returning a wrapped Rows object
func (op *Operation) Query(query string, args ...interface{}) *Rows {
	if op.Err() != nil {
		return &Rows{nil, op}
	}

	if op.Dbg {
		log.Printf("DEBUG - Querying: %s, %#v", query, args)
	}

	var r, err = op.q.Query(query, args...)
	op.SetErr(err)
	return &Rows{r, op}
}

// Exec wraps sql's DB.Exec, returning a wrapped Result
func (op *Operation) Exec(query string, args ...interface{}) *Result {
	if op.Err() != nil {
		return &Result{nil, op}
	}

	if op.Dbg {
		log.Printf("DEBUG - Executing: %s, %#v", query, args)
	}

	var r, err = op.q.Exec(query, args...)
	op.SetErr(err)
	return &Result{r, op}
}

// Prepare wrap's sql's DB.Prepare, returning a wrapped Stmt
func (op *Operation) Prepare(query string) *Stmt {
	if op.Err() != nil {
		return &Stmt{nil, op}
	}

	if op.Dbg {
		log.Printf("DEBUG - Preparing: %s", query)
	}

	var st, err = op.q.Prepare(query)
	op.SetErr(err)
	return &Stmt{st, op}
}

// Reset clears the error if any is present
func (op *Operation) Reset() {
	op.err = nil
}

// BeginTransaction wraps sql's Begin and uses a wrapped sql.Tx to dispatch
// Query, Exec, and Prepare calls.  When the transaction is complete, instead
// of manually rolling back or committing, simply call op.EndTransaction() and
// it will rollback / commit based on the error state.  If you need to force a
// rollback, set an error manually with Operation.SetErr().
//
// If a transaction is started while one is already in progress, the operation
// gets into an error state (i.e., nested transactions are not supported).
func (op *Operation) BeginTransaction() {
	if op.Err() != nil {
		return
	}

	if op.tx != nil {
		op.SetErr(errors.New("cannot nest transactions"))
		return
	}

	var tx, err = op.parent.db.Begin()
	op.tx = tx
	op.SetErr(err)
	op.q = tx
}

// EndTransaction commits the transaction if no errors occurred, or rolls back
// if there was an error
func (op *Operation) EndTransaction() {
	// If there was never a transaction due to errors, this could happen and we
	// don't want a panic
	if op.tx == nil {
		return
	}

	if op.Err() != nil {
		op.tx.Rollback()
	} else {
		op.SetErr(op.tx.Commit())
	}

	op.tx = nil
	op.q = op.parent.db
}

// Table creates an OperationTable tied to the given table name and reflecting
// on obj's type to auto-build certain SQL statements
func (op *Operation) Table(tableName string, obj interface{}) *OperationTable {
	var mt = &MagicTable{Object: obj, Name: tableName}
	mt.Configure(nil)
	return &OperationTable{op: op, t: mt}
}

// Save wraps Operation.Table() and Table.Save().  It creates an INSERT or
// UPDATE statement for the given object based on whether its primary key is
// zero.  Stores any errors the database returns, and fails if obj isn't tagged
// with a primary key field.
func (op *Operation) Save(tableName string, obj interface{}) *Result {
	var emptyResult = &Result{nil, op}

	if op.Err() != nil {
		return emptyResult
	}

	var ot = op.Table(tableName, obj)
	if ot.t.primaryKey == nil {
		op.SetErr(fmt.Errorf("no primary key tagged for structure %s", ot.t.RType.Name()))
		return emptyResult
	}

	return ot.Save(obj)
}

// Select wraps Operation.Table() and Table.Select().  It creates a Select
// object for further refining.
func (op *Operation) Select(tableName string, obj interface{}) Select {
	return op.Table(tableName, obj).Select()
}
<commit_msg>Add note about Operation.Prepare leaking resources<commit_after>package magicsql

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
)

// Querier defines an interface for top-level sql types that can run SQL and
// prepare statements
type Querier interface {
	Query(string, ...interface{}) (*sql.Rows, error)
	Exec(string, ...interface{}) (sql.Result, error)
	Prepare(string) (*sql.Stmt, error)
}

// Operation represents a short-lived single-purpose combination of database
// calls.  On the first failure, its internal error is set, which all
// "children" (statements, transactions, etc) will see.  All children will
// refuse to perform any functions once an error has occurred, making it safe
// to perform a chain of related database calls and only check for an error
// when it makes sense.
//
// When a transaction is started, the operation will route all database calls
// through the transaction instead of the global database handler.  At this
// time, only one transaction at a time is supported (i.e., no nesting
// transactions).
type Operation struct {
	parent *DB
	err    error
	tx     *sql.Tx
	q      Querier
	Dbg    bool
}

// NewOperation creates an operation in its default state: its parent is the
// passed-in DB instance, and it defaults to using direct database calls until
// a transaction is started.
func NewOperation(db *DB) *Operation {
	var o = &Operation{parent: db}
	o.q = o.parent.db

	return o
}

// Err returns the *first* error which occurred on any database call owned by the Operation
func (op *Operation) Err() error {
	return op.err
}

// SetErr tells the Operation to stop handling any more queries.  It shouldn't
// usually be called directly, but it can be if you need to tell the object
// "here's a thing that may be an error; don't do any more work if it is".
func (op *Operation) SetErr(err error) {
	if op.Err() != nil {
		return
	}

	op.err = err
}

// Query wraps sql's Query, returning a wrapped Rows object
func (op *Operation) Query(query string, args ...interface{}) *Rows {
	if op.Err() != nil {
		return &Rows{nil, op}
	}

	if op.Dbg {
		log.Printf("DEBUG - Querying: %s, %#v", query, args)
	}

	var r, err = op.q.Query(query, args...)
	op.SetErr(err)
	return &Rows{r, op}
}

// Exec wraps sql's DB.Exec, returning a wrapped Result
func (op *Operation) Exec(query string, args ...interface{}) *Result {
	if op.Err() != nil {
		return &Result{nil, op}
	}

	if op.Dbg {
		log.Printf("DEBUG - Executing: %s, %#v", query, args)
	}

	var r, err = op.q.Exec(query, args...)
	op.SetErr(err)
	return &Result{r, op}
}

// Prepare wrap's sql's DB.Prepare, returning a wrapped Stmt.  The statement
// must be closed by the caller or eventually MySQL will run out of prepared
// statements.
func (op *Operation) Prepare(query string) *Stmt {
	if op.Err() != nil {
		return &Stmt{nil, op}
	}

	if op.Dbg {
		log.Printf("DEBUG - Preparing: %s", query)
	}

	var st, err = op.q.Prepare(query)
	op.SetErr(err)
	return &Stmt{st, op}
}

// Reset clears the error if any is present
func (op *Operation) Reset() {
	op.err = nil
}

// BeginTransaction wraps sql's Begin and uses a wrapped sql.Tx to dispatch
// Query, Exec, and Prepare calls.  When the transaction is complete, instead
// of manually rolling back or committing, simply call op.EndTransaction() and
// it will rollback / commit based on the error state.  If you need to force a
// rollback, set an error manually with Operation.SetErr().
//
// If a transaction is started while one is already in progress, the operation
// gets into an error state (i.e., nested transactions are not supported).
func (op *Operation) BeginTransaction() {
	if op.Err() != nil {
		return
	}

	if op.tx != nil {
		op.SetErr(errors.New("cannot nest transactions"))
		return
	}

	var tx, err = op.parent.db.Begin()
	op.tx = tx
	op.SetErr(err)
	op.q = tx
}

// EndTransaction commits the transaction if no errors occurred, or rolls back
// if there was an error
func (op *Operation) EndTransaction() {
	// If there was never a transaction due to errors, this could happen and we
	// don't want a panic
	if op.tx == nil {
		return
	}

	if op.Err() != nil {
		op.tx.Rollback()
	} else {
		op.SetErr(op.tx.Commit())
	}

	op.tx = nil
	op.q = op.parent.db
}

// Table creates an OperationTable tied to the given table name and reflecting
// on obj's type to auto-build certain SQL statements
func (op *Operation) Table(tableName string, obj interface{}) *OperationTable {
	var mt = &MagicTable{Object: obj, Name: tableName}
	mt.Configure(nil)
	return &OperationTable{op: op, t: mt}
}

// Save wraps Operation.Table() and Table.Save().  It creates an INSERT or
// UPDATE statement for the given object based on whether its primary key is
// zero.  Stores any errors the database returns, and fails if obj isn't tagged
// with a primary key field.
func (op *Operation) Save(tableName string, obj interface{}) *Result {
	var emptyResult = &Result{nil, op}

	if op.Err() != nil {
		return emptyResult
	}

	var ot = op.Table(tableName, obj)
	if ot.t.primaryKey == nil {
		op.SetErr(fmt.Errorf("no primary key tagged for structure %s", ot.t.RType.Name()))
		return emptyResult
	}

	return ot.Save(obj)
}

// Select wraps Operation.Table() and Table.Select().  It creates a Select
// object for further refining.
func (op *Operation) Select(tableName string, obj interface{}) Select {
	return op.Table(tableName, obj).Select()
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"os"
	"syscall"
	"testing"
)

func TestBadEnv(t *testing.T) {
	buff := new(bytes.Buffer)
	stderr = buff
	exit = func(code int) {}
	main()
	os.Setenv("COFFEE_SHOP_CLOSE_TIME", "1")
	main()
	os.Setenv("COFFEE_SHOP_SHUTDOWN", "1")
	main()
	os.Setenv("COFFEE_SHOP_CUSTOMERS", "1")
	main()
}

func TestStore(t *testing.T) {
	buff := new(bytes.Buffer)
	stderr = buff
	exit = func(code int) {}
	os.Setenv("COFFEE_SHOP_CLOSE_TIME", "1")
	os.Setenv("COFFEE_SHOP_SHUTDOWN", "0")
	os.Setenv("COFFEE_SHOP_CUSTOMERS", "1")
	os.Setenv("COFFEE_SHOP_BARISTAS", "1")
	main()
}

func TestSignalStore(t *testing.T) {
	buff := new(bytes.Buffer)
	stderr = buff
	exit = func(code int) {}
	stop = make(chan os.Signal, 1)
	os.Setenv("COFFEE_SHOP_CLOSE_TIME", "1")
	os.Setenv("COFFEE_SHOP_SHUTDOWN", "1")
	os.Setenv("COFFEE_SHOP_CUSTOMERS", "1")
	os.Setenv("COFFEE_SHOP_BARISTAS", "1")
	stop <- syscall.SIGINT
	main()
}

func TestCantOpenStore(t *testing.T) {
	buff := new(bytes.Buffer)
	stderr = buff
	exit = func(code int) {}
	os.Setenv("COFFEE_SHOP_CLOSE_TIME", "0")
	os.Setenv("COFFEE_SHOP_SHUTDOWN", "0")
	os.Setenv("COFFEE_SHOP_CUSTOMERS", "0")
	os.Setenv("COFFEE_SHOP_BARISTAS", "0")
	main()
}
<commit_msg>make assertions for TestBadEnv<commit_after>package main

import (
	"bytes"
	"os"
	"strings"
	"syscall"
	"testing"
)

func TestBadEnv(t *testing.T) {
	buff := new(bytes.Buffer)
	stderr = buff
	var code int
	exit = func(c int) {
		code = c
	}
	main()
	if code != 1 {
		t.Error("expected exit status 1")
	}
	if !strings.Contains(buff.String(), "Bad COFFEE_SHOP_CLOSE_TIME") {
		t.Error("Didn't provide close time, but didn't error")
	}
	code = 0
	os.Setenv("COFFEE_SHOP_CLOSE_TIME", "1")
	main()
	if code != 1 {
		t.Error("expected exit status 1")
	}
	if !strings.Contains(buff.String(), "Bad COFFEE_SHOP_SHUTDOWN") {
		t.Error("Didn't provide shutdown time, but didn't error")
	}
	code = 0
	os.Setenv("COFFEE_SHOP_SHUTDOWN", "1")
	main()
	if code != 1 {
		t.Error("expected exit status 1")
	}
	if !strings.Contains(buff.String(), "Bad COFFEE_SHOP_CUSTOMERS") {
		t.Error("Didn't provide number of customers, but didn't error")
	}
	code = 0
	os.Setenv("COFFEE_SHOP_CUSTOMERS", "1")
	main()
	if code != 1 {
		t.Error("expected exit status 1")
	}
	if !strings.Contains(buff.String(), "Bad COFFEE_SHOP_BARISTAS") {
		t.Error("Didn't provide number of baristas, but didn't error")
	}
}

func TestStore(t *testing.T) {
	buff := new(bytes.Buffer)
	stderr = buff
	exit = func(code int) {}
	os.Setenv("COFFEE_SHOP_CLOSE_TIME", "1")
	os.Setenv("COFFEE_SHOP_SHUTDOWN", "0")
	os.Setenv("COFFEE_SHOP_CUSTOMERS", "1")
	os.Setenv("COFFEE_SHOP_BARISTAS", "1")
	main()
}

func TestSignalStore(t *testing.T) {
	buff := new(bytes.Buffer)
	stderr = buff
	exit = func(code int) {}
	stop = make(chan os.Signal, 1)
	os.Setenv("COFFEE_SHOP_CLOSE_TIME", "1")
	os.Setenv("COFFEE_SHOP_SHUTDOWN", "1")
	os.Setenv("COFFEE_SHOP_CUSTOMERS", "1")
	os.Setenv("COFFEE_SHOP_BARISTAS", "1")
	stop <- syscall.SIGINT
	main()
}

func TestCantOpenStore(t *testing.T) {
	buff := new(bytes.Buffer)
	stderr = buff
	exit = func(code int) {}
	os.Setenv("COFFEE_SHOP_CLOSE_TIME", "0")
	os.Setenv("COFFEE_SHOP_SHUTDOWN", "0")
	os.Setenv("COFFEE_SHOP_CUSTOMERS", "0")
	os.Setenv("COFFEE_SHOP_BARISTAS", "0")
	main()
}
<|endoftext|>


<commit_before>package main

import (
	"testing"

	"github.com/Liamraystanley/marill/scraper"
)

func TestFetch(t *testing.T) {
	cases := []struct {
		in   string
		want bool
	}{
		{"https://liamstanley.io", false},
		{"https://google.com/", false},
		{"htps://google.com", true},
		{"http://liamstanley.io", false},
		{"http://some-domains-that-doesnt-exist.com", true},
		{"https://some-domains-that-doesnt-exist.com", true},
	}

	for _, c := range cases {
		got := scraper.Crawl(c.in, "")

		if got.Error != nil && !c.want {
			t.Errorf("scraper.Crawl(%q) == %q, wanted error: %v", c.in, got.Error, c.want)
		}

		if got.Error == nil && c.want && got.Code == 200 {
			t.Errorf("scraper.Crawl(%s) == %q (%q), though no errors", c.in, got.Error, got)
		}
	}

	return
}
<commit_msg>more tests<commit_after>package main

import (
	"testing"

	"github.com/Liamraystanley/marill/scraper"
)

func TestFetch(t *testing.T) {
	cases := []struct {
		in   string
		inx  string // extra -- e.g. ip
		want bool
	}{
		{"https://liamstanley.io", "", false},
		{"https://liamstanley.io", "0.0.0.0", false},
		{"https://liamstanley.io", "000.000.00.0", false},
		{"https://liamstanley.io", "111.1111.11.1", false},
		{"https://liamstanley.io", "1.1.1.1.", false},
		{"https://google.com/", "", false},
		{"htps://google.com", "", true},
		{"http://liamstanley.io", "", false},
		{"http://liamstanley.io", "0.0.0.0", false},
		{"http://liamstanley.io", "000.000.00.0", false},
		{"http://liamstanley.io", "111.1111.11.1", false},
		{"http://liamstanley.io", "1.1.1.1.", false},
		{"http://some-domains-that-doesnt-exist.com", "", true},
		{"https://some-domains-that-doesnt-exist.com", "", true},
	}

	for _, c := range cases {
		got := scraper.FetchURL(c.in, "")

		if got.Error != nil && !c.want {
			t.Errorf("scraper.fetchURL(%q, %q) == %q, wanted error: %v", c.in, c.inx, got.Error, c.want)
		}

		if got.Error == nil && c.want && got.Code == 200 {
			t.Errorf("scraper.fetchURL(%q, %q) == %q (%#v), though no errors", c.in, c.inx, got.Error, got)
		}
	}

	return
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"fmt"
	"github.com/op/go-logging"
	. "github.com/smartystreets/goconvey/convey"
	"image"
	_ "image/png"
	"testing"
)

const (
	testUser = "clone1018"
)

type SilentWriter struct {
}

func (w SilentWriter) Write(p []byte) (n int, err error) {
	return 0, nil
}

func hashRender(render image.Image) string {
	// md5 hash its pixels
	hasher := md5.New()
	hasher.Write(render.(*image.NRGBA).Pix)
	md5Hash := fmt.Sprintf("%x", hasher.Sum(nil))

	// And return the Hash
	return md5Hash
}

func TestSetup(t *testing.T) {
	logBackend := logging.NewLogBackend(SilentWriter{}, "", 0)
	stats = MakeStatsCollector()
	setupConfig()
	setupLog(logBackend)
	setupCache()
}

func TestRenders(t *testing.T) {
	Convey("GetHead should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetHead(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "257972f3af185ce1f92200156d3bfe97")
	})

	Convey("GetHelm should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetHelm(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "257972f3af185ce1f92200156d3bfe97")
	})

	Convey("GetCube should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetCube(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "a579af057377c35cbd4e53ed6bc5f03e")
	})

	Convey("GetBust should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetBust(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "1a4ba8567619350883923bde97626ae6")
	})

	Convey("GetBody should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetBody(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "987cc81d031386f429b12da2cd5c1ff2")
	})

	Convey("GetArmorBust should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetArmorBust(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "1a4ba8567619350883923bde97626ae6")
	})

	Convey("GetArmorBody should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetArmorBody(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "987cc81d031386f429b12da2cd5c1ff2")
	})
}

func BenchmarkSetup(b *testing.B) {
	logBackend := logging.NewLogBackend(SilentWriter{}, "", 0)
	stats = MakeStatsCollector()
	setupConfig()
	setupLog(logBackend)
	setupCache()
}

func BenchmarkGetHead(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetHead(20)
	}
}

func BenchmarkGetHelm(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetHelm(20)
	}
}

func BenchmarkGetCube(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetCube(20)
	}
}

func BenchmarkGetBust(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetBust(20)
	}
}

func BenchmarkGetBody(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetBody(20)
	}
}

func BenchmarkGetArmorBust(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetArmorBust(20)
	}
}

func BenchmarkGetArmorBody(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetArmorBody(20)
	}
}
<commit_msg>Hotfix test fix<commit_after>package main

import (
	"crypto/md5"
	"fmt"
	"github.com/op/go-logging"
	. "github.com/smartystreets/goconvey/convey"
	"image"
	_ "image/png"
	"testing"
)

const (
	testUser = "clone1018"
)

type SilentWriter struct {
}

func (w SilentWriter) Write(p []byte) (n int, err error) {
	return 0, nil
}

func hashRender(render image.Image) string {
	// md5 hash its pixels
	hasher := md5.New()
	hasher.Write(render.(*image.NRGBA).Pix)
	md5Hash := fmt.Sprintf("%x", hasher.Sum(nil))

	// And return the Hash
	return md5Hash
}

func TestSetup(t *testing.T) {
	logBackend := logging.NewLogBackend(SilentWriter{}, "", 0)
	stats = MakeStatsCollector()
	setupConfig()
	setupLog(logBackend)
	setupCache()
}

func TestRenders(t *testing.T) {
	Convey("GetHead should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetHead(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "257972f3af185ce1f92200156d3bfe97")
	})

	Convey("GetHelm should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetHelm(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "257972f3af185ce1f92200156d3bfe97")
	})

	Convey("GetCube should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetCube(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "a253bb68f5ed938eb235ff1e3807940c")
	})

	Convey("GetBust should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetBust(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "1a4ba8567619350883923bde97626ae6")
	})

	Convey("GetBody should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetBody(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "987cc81d031386f429b12da2cd5c1ff2")
	})

	Convey("GetArmorBust should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetArmorBust(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "1a4ba8567619350883923bde97626ae6")
	})

	Convey("GetArmorBody should return a valid image", t, func() {
		skin := fetchSkin(testUser)
		err := skin.GetArmorBody(20)

		So(skin.Processed, ShouldNotBeNil)
		So(err, ShouldBeNil)

		hash := hashRender(skin.Processed)
		So(hash, ShouldEqual, "987cc81d031386f429b12da2cd5c1ff2")
	})
}

func BenchmarkSetup(b *testing.B) {
	logBackend := logging.NewLogBackend(SilentWriter{}, "", 0)
	stats = MakeStatsCollector()
	setupConfig()
	setupLog(logBackend)
	setupCache()
}

func BenchmarkGetHead(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetHead(20)
	}
}

func BenchmarkGetHelm(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetHelm(20)
	}
}

func BenchmarkGetCube(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetCube(20)
	}
}

func BenchmarkGetBust(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetBust(20)
	}
}

func BenchmarkGetBody(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetBody(20)
	}
}

func BenchmarkGetArmorBust(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetArmorBust(20)
	}
}

func BenchmarkGetArmorBody(b *testing.B) {
	skin := fetchSkin(testUser)
	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		skin.GetArmorBody(20)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"reflect"
	"sort"
	"strings"
	"testing"
	"time"
)

func TestScanToChannel(t *testing.T) {
	stdin := `1
2`

	from := strings.NewReader(stdin)
	to := make(chan string, 2)
	cancel := make(chan struct{})

	scanToChannel(from, to, cancel)

	lines := readAndSortBlocking(to, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"1", "2"}) != true {
		t.Errorf("result wasn't ['1', '2'], it was %v", lines)
	}
}

func TestReadCmd(t *testing.T) {
	cmdString := `echo -e "1\n2"`
	o := make(chan string, 2)
	cancel := make(chan struct{})

	readCmd(cmdString, o, cancel)

	lines := readAndSortBlocking(o, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"1", "2"}) != true {
		t.Errorf("result wasn't ['1', '2'], it was %v", lines)
	}
}

type mockUtils struct {
	i io.Reader
}

func (m mockUtils) scanStdinToChannel(i chan string, cancel chan struct{}) {
	scanToChannel(m.i, i, cancel)
}

func TestDiff(t *testing.T) {
	stdout := make(chan string)
	stdin := `1
2
3
4`
	reader := strings.NewReader(stdin)

	go diff(`echo -e "1\n2"`, defaultTimeout(), defaultTimeout(), stdout, mockUtils{reader})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"3", "4"}) != true {
		t.Errorf("result wasn't ['3', '4'], it was %v", lines)
	}
}

func TestDiffWhenInputTimesOut(t *testing.T) {
	stdout := make(chan string)

	reader := func() io.Reader {
		reader, writer := io.Pipe()
		go func() {
			io.WriteString(writer, "1\n")
			io.WriteString(writer, "3\n")
			io.WriteString(writer, "3\n")
			io.WriteString(writer, "3\n")
			io.WriteString(writer, "1\n")
			io.WriteString(writer, "2\n")
			io.WriteString(writer, "4\n")
			time.Sleep(101 * time.Millisecond)
			io.WriteString(writer, "5\n")
			writer.Close()
		}()
		return reader
	}()

	go diff(`echo -e "1\n2"`, defaultTimeout(), defaultTimeout(), stdout, mockUtils{reader})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	sort.Strings(lines) // order is not deterministic
	if reflect.DeepEqual(lines, []string{"3", "3", "3", "4"}) != true {
		t.Errorf("result wasn't ['3', '3', '3', '4'], it was %v", lines)
	}
}

func TestDiffWhenOutputTimesOut(t *testing.T) {
	stdout := make(chan string)
	stdin := `1
2
3
4
5`
	reader := strings.NewReader(stdin)

	go diff(`echo -e "1\n2" && sleep 1 && echo -e "3\n4"`, defaultTimeout(), defaultTimeout(), stdout, mockUtils{reader})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"3", "4", "5"}) != true {
		t.Errorf("result wasn't ['3', '4', '5'], it was %v", lines)
	}
}

func TestDiffWhenDelaysAddUpToTimeoutSeparatelyButDoesntTimeout(t *testing.T) {
	stdout := make(chan string)
	stdin := `1
2
3
4
5
6`
	reader := strings.NewReader(stdin)

	go diff(`echo "1" && sleep .1 && echo "2" && sleep .1 && echo "3" && sleep .1 && echo "4" && sleep .1 && echo "ten"`,
		defaultTimeout(), timeout{firstTime: 200 * time.Millisecond, time: 200 * time.Millisecond}, stdout, mockUtils{reader})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"5", "6"}) != true {
		t.Errorf("result wasn't ['5', '6'], it was %v", lines)
	}
}

func readAndSortBlocking(c chan string, timeout time.Duration) []string {
	lines := []string{}
loop:
	for {
		select {
		case line, ok := <-c:
			if !ok {
				break loop
			}
			lines = append(lines, line)
		case <-time.After(timeout):
			break loop
		}
	}
	sort.Strings(lines) // order is not deterministic

	return lines
}

func TestEmptyCommand(t *testing.T) {
	stdout := make(chan string)

	stdin := `1
2
3`
	reader := strings.NewReader(stdin)
	go diff(`echo ""`, defaultTimeout(), defaultTimeout(), stdout, mockUtils{reader})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"1", "2", "3"}) != true {
		t.Errorf("result wasn't ['1', '2', '3'], it was %v", lines)
	}
}

func TestEmptyStdin(t *testing.T) {
	stdout := make(chan string)

	go diff(`echo "1\n2\n3"`, defaultTimeout(), defaultTimeout(), stdout, mockUtils{strings.NewReader(``)})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	if reflect.DeepEqual(lines, []string{}) != true {
		t.Errorf("result wasn't [], it was %v", lines)
	}
}

func defaultTimeout() timeout {
	return timeout{firstTime: 100 * time.Millisecond, time: 100 * time.Millisecond}
}
<commit_msg>Implements aux function to test using ext. procs.<commit_after>package main

import (
	"io"
	"log"
	"os/exec"
	"reflect"
	"sort"
	"strings"
	"testing"
	"time"
)

func TestScanToChannel(t *testing.T) {
	stdin := `1
2`

	from := strings.NewReader(stdin)
	to := make(chan string, 2)
	cancel := make(chan struct{})

	scanToChannel(from, to, cancel)

	lines := readAndSortBlocking(to, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"1", "2"}) != true {
		t.Errorf("result wasn't ['1', '2'], it was %v", lines)
	}
}

func TestReadCmd(t *testing.T) {
	cmdString := `echo -e "1\n2"`
	o := make(chan string, 2)
	cancel := make(chan struct{})

	readCmd(cmdString, o, cancel)

	lines := readAndSortBlocking(o, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"1", "2"}) != true {
		t.Errorf("result wasn't ['1', '2'], it was %v", lines)
	}
}

type mockUtils struct {
	i io.Reader
}

func (m mockUtils) scanStdinToChannel(i chan string, cancel chan struct{}) {
	scanToChannel(m.i, i, cancel)
}

func TestDiff(t *testing.T) {
	stdout := make(chan string)
	reader := cmdToReader(`echo -e "1\n2\n3\n4"`)

	go diff(`echo -e "1\n2"`, defaultTimeout(), defaultTimeout(), stdout, mockUtils{reader})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"3", "4"}) != true {
		t.Errorf("result wasn't ['3', '4'], it was %v", lines)
	}
}

func TestDiffWhenInputTimesOut(t *testing.T) {
	stdout := make(chan string)

	reader := cmdToReader(`echo -e "1\n3\n3\n3\n1\n2\n4" && sleep .101 && echo "5"`)

	go diff(`echo -e "1\n2"`, defaultTimeout(), defaultTimeout(), stdout, mockUtils{reader})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	sort.Strings(lines) // order is not deterministic
	if reflect.DeepEqual(lines, []string{"3", "3", "3", "4"}) != true {
		t.Errorf("result wasn't ['3', '3', '3', '4'], it was %v", lines)
	}
}

func TestDiffWhenOutputTimesOut(t *testing.T) {
	stdout := make(chan string)
	reader := cmdToReader(`echo -e "1\n2\n3\n4\n5"`)

	go diff(`echo -e "1\n2" && sleep 1 && echo -e "3\n4"`, defaultTimeout(), defaultTimeout(), stdout, mockUtils{reader})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"3", "4", "5"}) != true {
		t.Errorf("result wasn't ['3', '4', '5'], it was %v", lines)
	}
}

func TestDiffWhenDelaysAddUpToTimeoutSeparatelyButDoesntTimeout(t *testing.T) {
	stdout := make(chan string)
	reader := cmdToReader(`echo -e "1\n2\n3\n4\n5\n6"`)

	go diff(`echo "1" && sleep .1 && echo "2" && sleep .1 && echo "3" && sleep .1 && echo "4" && sleep .1 && echo "ten"`,
		defaultTimeout(), timeout{firstTime: 200 * time.Millisecond, time: 200 * time.Millisecond}, stdout, mockUtils{reader})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"5", "6"}) != true {
		t.Errorf("result wasn't ['5', '6'], it was %v", lines)
	}
}

func readAndSortBlocking(c chan string, timeout time.Duration) []string {
	lines := []string{}
loop:
	for {
		select {
		case line, ok := <-c:
			if !ok {
				break loop
			}
			lines = append(lines, line)
		case <-time.After(timeout):
			break loop
		}
	}
	sort.Strings(lines) // order is not deterministic

	return lines
}

func TestEmptyCommand(t *testing.T) {
	stdout := make(chan string)

	reader := cmdToReader(`echo -e "1\n2\n3"`)
	go diff(`echo ""`, defaultTimeout(), defaultTimeout(), stdout, mockUtils{reader})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	if reflect.DeepEqual(lines, []string{"1", "2", "3"}) != true {
		t.Errorf("result wasn't ['1', '2', '3'], it was %v", lines)
	}
}

func TestEmptyStdin(t *testing.T) {
	stdout := make(chan string)

	go diff(`echo "1\n2\n3"`, defaultTimeout(), defaultTimeout(), stdout, mockUtils{strings.NewReader(``)})

	lines := readAndSortBlocking(stdout, 1*time.Second)

	if reflect.DeepEqual(lines, []string{}) != true {
		t.Errorf("result wasn't [], it was %v", lines)
	}
}

func defaultTimeout() timeout {
	return timeout{firstTime: 100 * time.Millisecond, time: 100 * time.Millisecond}
}

func cmdToReader(cmdString string) io.Reader {
	cmd := exec.Command("bash", "-c", cmdString)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		log.Fatal(err)
	}

	if err := cmd.Start(); err != nil {
		log.Fatal(err)
	}

	return stdout
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/olekukonko/tablewriter"
	"github.com/pinzolo/dbmodel"
)

var cmdIndex = &Command{
	Run:       runIndex,
	UsageLine: "index ",
	Short:     "Enumerate all table names in schema.",
	Long: `Enumerate table names.

Options:
    -c 'config gile', --config 'config file'
        use config file instead of default config file(tablarian.config)
        if 'config file' starts with '@', it is treated as absolute file path.

    -C, --no-comment
	    Not print table comment. (default: false)
	`,
}

var withoutTableComment = false

func init() {
	cmdIndex.Flag.StringVar(&configFile, "config", "tablarian.config", "Config file path")
	cmdIndex.Flag.StringVar(&configFile, "c", "tablarian.config", "Config file path")
	cmdIndex.Flag.BoolVar(&withoutTableComment, "no-comment", false, "Without table comment")
	cmdIndex.Flag.BoolVar(&withoutTableComment, "C", false, "Without table comment")
}

// runIndex executes index command and return exit code.
func runIndex(args []string) int {
	cfg, err := loadConfig(configFile)
	if err != nil {
		fmt.Fprintln(o.err, err)
		return 1
	}
	db := dbClientFor(cfg)
	db.Connect()
	defer db.Disconnect()

	tables, err := db.AllTableNames(cfg.Schema)
	if err != nil {
		fmt.Fprintln(o.err, err)
		return 1
	}

	printTableNames(tables)
	return 0
}

func printTableNames(tables []*dbmodel.Table) {
	for _, line := range tableNamesLines(tables) {
		fmt.Fprintln(o.out, line)
	}
}

func tableNamesLines(tables []*dbmodel.Table) []string {
	buf := &bytes.Buffer{}
	w := tablewriter.NewWriter(buf)
	w.SetBorder(false)
	w.SetColumnSeparator("")
	w.SetAutoWrapText(false)
	for _, tbl := range tables {
		data := []string{tbl.Name()}
		if !withoutTableComment {
			data = append(data, tbl.Comment())
		}
		w.Append(data)
	}
	w.Render()
	return trimEachLines(buf.String())
}

func trimEachLines(tableString string) []string {
	lines := strings.Split(tableString, "\n")
	newLines := make([]string, 0, len(lines))
	for _, line := range lines {
		if l := strings.TrimSpace(line); len(l) > 0 {
			newLines = append(newLines, l)
		}
	}
	return newLines
}
<commit_msg>Fix usage [ci skip]<commit_after>package main

import (
	"bytes"
	"fmt"
	"strings"

	"github.com/olekukonko/tablewriter"
	"github.com/pinzolo/dbmodel"
)

var cmdIndex = &Command{
	Run:       runIndex,
	UsageLine: "index ",
	Short:     "Enumerate all table names in schema.",
	Long: `Enumerate table names in schema.

Options:
    -c 'config gile', --config 'config file'
        use config file instead of default config file(tablarian.config)
        if 'config file' starts with '@', it is treated as absolute file path.

    -C, --no-comment
	    Not print table comment. (default: false)
	`,
}

var withoutTableComment = false

func init() {
	cmdIndex.Flag.StringVar(&configFile, "config", "tablarian.config", "Config file path")
	cmdIndex.Flag.StringVar(&configFile, "c", "tablarian.config", "Config file path")
	cmdIndex.Flag.BoolVar(&withoutTableComment, "no-comment", false, "Without table comment")
	cmdIndex.Flag.BoolVar(&withoutTableComment, "C", false, "Without table comment")
}

// runIndex executes index command and return exit code.
func runIndex(args []string) int {
	cfg, err := loadConfig(configFile)
	if err != nil {
		fmt.Fprintln(o.err, err)
		return 1
	}
	db := dbClientFor(cfg)
	db.Connect()
	defer db.Disconnect()

	tables, err := db.AllTableNames(cfg.Schema)
	if err != nil {
		fmt.Fprintln(o.err, err)
		return 1
	}

	printTableNames(tables)
	return 0
}

func printTableNames(tables []*dbmodel.Table) {
	for _, line := range tableNamesLines(tables) {
		fmt.Fprintln(o.out, line)
	}
}

func tableNamesLines(tables []*dbmodel.Table) []string {
	buf := &bytes.Buffer{}
	w := tablewriter.NewWriter(buf)
	w.SetBorder(false)
	w.SetColumnSeparator("")
	w.SetAutoWrapText(false)
	for _, tbl := range tables {
		data := []string{tbl.Name()}
		if !withoutTableComment {
			data = append(data, tbl.Comment())
		}
		w.Append(data)
	}
	w.Render()
	return trimEachLines(buf.String())
}

func trimEachLines(tableString string) []string {
	lines := strings.Split(tableString, "\n")
	newLines := make([]string, 0, len(lines))
	for _, line := range lines {
		if l := strings.TrimSpace(line); len(l) > 0 {
			newLines = append(newLines, l)
		}
	}
	return newLines
}
<|endoftext|>


<commit_before>package tiles

import (
	"bytes"
	"index/suffixarray"
	"sort"
	"sync"
)

const (
	zero = byte('\x00')
)

// TileIndex stores indexes values by tile.
// If a deep level of tile is added and a shallower one is requested, the values are aggregated up.
type TileIndex interface {
	TileRange(zmin, zmax int) <-chan Tile
	Values(t Tile) (vals []interface{})
	Add(t Tile, val ...interface{})
}

// NewTileIndex returns the default TileIndex
func NewTileIndex() TileIndex {
	return &KeysetIndex{}
}

// KeysetIndex is a TileIndex implementation that uses a sorted keyset.
// A trie would be more efficient, but KeysetIndex mirrors the range queries of boltdb which could be dropped in if the entire index won't fit in memory.
// KeysetIndex is thread safe
type KeysetIndex struct {
	// Implementation uses a sorted keyset.
	// A trie would be more efficient, but
	sorted bool
	keys   []qkey
	values [][]interface{}
	sync.RWMutex
}

// TileRange returns a channel of all tiles in the index in the zoom range
// If zmax is greater than the deepest tile level, the deepest tile level returns
// Acquires a readlock for duration of returned channel being open
func (idx *KeysetIndex) TileRange(zmin, zmax int) <-chan Tile {
	idx.sort()
	tiles := make(chan Tile, 1<<10)
	go func() {
		defer close(tiles)
		idx.RLock()
		defer idx.RUnlock()
		for i := 0; i < len(idx.keys)-1; i++ {
			qmax := idx.keys[i].qk.Level()
			for z := zmin; z <= zmax && z <= qmax; z++ {
				q := idx.keys[i].qk.Parent(z)
				n := idx.keys[i+1].qk
				if !n.HasParent(q) {
					tiles <- q.ToTile()
				}
			}
		}
		q := idx.keys[len(idx.keys)-1].qk
		for z := zmin; z <= zmax && z <= len(q); z++ {
			tiles <- q.Parent(z).ToTile()
		}
	}()
	return tiles
}

// Values returns a list of values aggregated under the requested tile
func (idx *KeysetIndex) Values(t Tile) (vals []interface{}) {
	idx.sort()
	idx.RLock()
	defer idx.RUnlock()
	qk := t.Quadkey()
	i := idx.search(qk)
	for i < len(idx.keys) {
		n := idx.keys[i]
		if n.qk == qk || n.qk.HasParent(qk) {
			vals = append(vals, idx.values[n.v]...)
		}
		i++
	}
	return
}

// Add adds a value, but will not be indexed
func (idx *KeysetIndex) Add(t Tile, val ...interface{}) {
	idx.Lock()
	defer idx.Unlock()
	idx.values = append(idx.values, val)
	qk := qkey{qk: t.Quadkey(), v: len(idx.values) - 1}
	idx.keys = append(idx.keys, qk)
	idx.sorted = false
}

// sorts the tiles, nothing happens if the sorted flag is set
func (idx *KeysetIndex) sort() {
	if !idx.sorted {
		idx.Lock()
		sort.Sort(byQk(idx.keys))
		idx.sorted = true
		idx.Unlock()
	}
}

func (idx *KeysetIndex) search(qk Quadkey) int {
	return sort.Search(len(idx.keys), func(i int) bool { return idx.keys[i].qk >= qk })
}

type qkey struct {
	qk Quadkey
	v  int
}

type byQk []qkey

func (q byQk) Len() int           { return len(q) }
func (q byQk) Swap(i, j int)      { q[i], q[j] = q[j], q[i] }
func (q byQk) Less(i, j int) bool { return q[i].qk < q[j].qk }

//SuffixIndex is a TileIndex that uses a suffixarray to lookup values
//It IS NOT currently safe for concurrent access.
type SuffixIndex struct {
	// \x00 joined string of keys for suffixarray
	indexed []byte
	index   *suffixarray.Index
	tiles   map[Quadkey][]interface{}
}

//NewSuffixIndex returns a new SuffixIndex
func NewSuffixIndex() *SuffixIndex {
	return &SuffixIndex{
		tiles: make(map[Quadkey][]interface{}),
	}
}

//TileRange returns all the tiles available in this index.
//It currently DOES NOT return unique values
func (idx *SuffixIndex) TileRange(zmin, zmax int) <-chan Tile {
	tiles := make(chan Tile, 1<<10)
	go func() {
		for k := range idx.tiles {
			for z := zmin; z <= zmax; z++ {
				tiles <- k[:z].ToTile()
			}
		}
	}()
	return tiles
}

//Values returns all the values aggregated under the given tile
func (idx *SuffixIndex) Values(t Tile) (vals []interface{}) {
	idx.sort()
	qk := t.Quadkey()
	keys := prefixes(idx.index, idx.indexed, []byte(qk))
	for _, k := range keys {
		qk := Quadkey(k)
		//fmt.Println(qk)
		vals = append(vals, idx.tiles[qk]...)
	}
	return
}

//Add adds a tile and values associated with it
func (idx *SuffixIndex) Add(t Tile, v ...interface{}) {
	// Set index to nil b/c adding invalidates index
	idx.index = nil
	qk := t.Quadkey()
	idx.tiles[qk] = append(idx.tiles[qk], v...)
}

func (idx *SuffixIndex) sort() {
	if idx.index == nil {
		keys := make([][]byte, len(idx.tiles))
		i := 0
		for k := range idx.tiles {
			keys[i] = []byte(k)
			i++
		}
		d := []byte{zero}
		b := bytes.Join(keys, d)                    //join w/ zeros
		idx.indexed = bytes.Join([][]byte{d, d}, b) //pad w/ zeros
		idx.index = suffixarray.New(idx.indexed)
	}
}

//prefixes assumes a \x00 delimited data with \x00 padding
func prefixes(idx *suffixarray.Index, data, q []byte) (keys [][]byte) {
	for _, i := range idx.Lookup(q, -1) {
		if data[i-1] == zero { //if previous byte is zero, it's a prefix
			var l int
			for l = i; data[l] != zero; l++ {
				// iterating until end of word
			}
			k := data[i:l]
			keys = append(keys, k)
		}
	}
	return
}
<commit_msg>index: suffix close tilerange<commit_after>package tiles

import (
	"bytes"
	"index/suffixarray"
	"sort"
	"sync"
)

const (
	zero = byte('\x00')
)

// TileIndex stores indexes values by tile.
// If a deep level of tile is added and a shallower one is requested, the values are aggregated up.
type TileIndex interface {
	TileRange(zmin, zmax int) <-chan Tile
	Values(t Tile) (vals []interface{})
	Add(t Tile, val ...interface{})
}

// NewTileIndex returns the default TileIndex
func NewTileIndex() TileIndex {
	return &KeysetIndex{}
}

// KeysetIndex is a TileIndex implementation that uses a sorted keyset.
// A trie would be more efficient, but KeysetIndex mirrors the range queries of boltdb which could be dropped in if the entire index won't fit in memory.
// KeysetIndex is thread safe
type KeysetIndex struct {
	// Implementation uses a sorted keyset.
	// A trie would be more efficient, but
	sorted bool
	keys   []qkey
	values [][]interface{}
	sync.RWMutex
}

// TileRange returns a channel of all tiles in the index in the zoom range
// If zmax is greater than the deepest tile level, the deepest tile level returns
// Acquires a readlock for duration of returned channel being open
func (idx *KeysetIndex) TileRange(zmin, zmax int) <-chan Tile {
	idx.sort()
	tiles := make(chan Tile, 1<<10)
	go func() {
		defer close(tiles)
		idx.RLock()
		defer idx.RUnlock()
		for i := 0; i < len(idx.keys)-1; i++ {
			qmax := idx.keys[i].qk.Level()
			for z := zmin; z <= zmax && z <= qmax; z++ {
				q := idx.keys[i].qk.Parent(z)
				n := idx.keys[i+1].qk
				if !n.HasParent(q) {
					tiles <- q.ToTile()
				}
			}
		}
		q := idx.keys[len(idx.keys)-1].qk
		for z := zmin; z <= zmax && z <= len(q); z++ {
			tiles <- q.Parent(z).ToTile()
		}
	}()
	return tiles
}

// Values returns a list of values aggregated under the requested tile
func (idx *KeysetIndex) Values(t Tile) (vals []interface{}) {
	idx.sort()
	idx.RLock()
	defer idx.RUnlock()
	qk := t.Quadkey()
	i := idx.search(qk)
	for i < len(idx.keys) {
		n := idx.keys[i]
		if n.qk == qk || n.qk.HasParent(qk) {
			vals = append(vals, idx.values[n.v]...)
		}
		i++
	}
	return
}

// Add adds a value, but will not be indexed
func (idx *KeysetIndex) Add(t Tile, val ...interface{}) {
	idx.Lock()
	defer idx.Unlock()
	idx.values = append(idx.values, val)
	qk := qkey{qk: t.Quadkey(), v: len(idx.values) - 1}
	idx.keys = append(idx.keys, qk)
	idx.sorted = false
}

// sorts the tiles, nothing happens if the sorted flag is set
func (idx *KeysetIndex) sort() {
	if !idx.sorted {
		idx.Lock()
		sort.Sort(byQk(idx.keys))
		idx.sorted = true
		idx.Unlock()
	}
}

func (idx *KeysetIndex) search(qk Quadkey) int {
	return sort.Search(len(idx.keys), func(i int) bool { return idx.keys[i].qk >= qk })
}

type qkey struct {
	qk Quadkey
	v  int
}

type byQk []qkey

func (q byQk) Len() int           { return len(q) }
func (q byQk) Swap(i, j int)      { q[i], q[j] = q[j], q[i] }
func (q byQk) Less(i, j int) bool { return q[i].qk < q[j].qk }

//SuffixIndex is a TileIndex that uses a suffixarray to lookup values
//It IS NOT currently safe for concurrent access.
type SuffixIndex struct {
	// \x00 joined string of keys for suffixarray
	indexed []byte
	index   *suffixarray.Index
	tiles   map[Quadkey][]interface{}
}

//NewSuffixIndex returns a new SuffixIndex
func NewSuffixIndex() *SuffixIndex {
	return &SuffixIndex{
		tiles: make(map[Quadkey][]interface{}),
	}
}

//TileRange returns all the tiles available in this index.
//It currently DOES NOT return unique values
func (idx *SuffixIndex) TileRange(zmin, zmax int) <-chan Tile {
	tiles := make(chan Tile, 1<<10)
	go func() {
		defer close(tiles)
		for k := range idx.tiles {
			for z := zmin; z <= zmax; z++ {
				tiles <- k[:z].ToTile()
			}
		}
	}()
	return tiles
}

//Values returns all the values aggregated under the given tile
func (idx *SuffixIndex) Values(t Tile) (vals []interface{}) {
	idx.sort()
	qk := t.Quadkey()
	keys := prefixes(idx.index, idx.indexed, []byte(qk))
	for _, k := range keys {
		qk := Quadkey(k)
		//fmt.Println(qk)
		vals = append(vals, idx.tiles[qk]...)
	}
	return
}

//Add adds a tile and values associated with it
func (idx *SuffixIndex) Add(t Tile, v ...interface{}) {
	// Set index to nil b/c adding invalidates index
	idx.index = nil
	qk := t.Quadkey()
	idx.tiles[qk] = append(idx.tiles[qk], v...)
}

func (idx *SuffixIndex) sort() {
	if idx.index == nil {
		keys := make([][]byte, len(idx.tiles))
		i := 0
		for k := range idx.tiles {
			keys[i] = []byte(k)
			i++
		}
		d := []byte{zero}
		b := bytes.Join(keys, d)                    //join w/ zeros
		idx.indexed = bytes.Join([][]byte{d, d}, b) //pad w/ zeros
		idx.index = suffixarray.New(idx.indexed)
	}
}

//prefixes assumes a \x00 delimited data with \x00 padding
func prefixes(idx *suffixarray.Index, data, q []byte) (keys [][]byte) {
	for _, i := range idx.Lookup(q, -1) {
		if data[i-1] == zero { //if previous byte is zero, it's a prefix
			var l int
			for l = i; data[l] != zero; l++ {
				// iterating until end of word
			}
			k := data[i:l]
			keys = append(keys, k)
		}
	}
	return
}
<|endoftext|>


<commit_before>package main

// https://cloud.google.com/appengine/docs/go/tools/remoteapi
import _ "appengine/remote_api"

import (
	"encoding/json"
	"html/template"
	"net/http"

	"appengine"
	"appengine/datastore"
	"appengine/user"
)

var (
	templateIndex = template.Must(template.ParseFiles(
		"templates/index.html",
	))
	template404 = template.Must(template.ParseFiles(
		"templates/404.html",
	))
	frequencyResponses = map[int]string{
		1:  `["once a week", "week"]`,
		4:  `["every two days", "two-day"]`,
		7:  `["once a day", "day"]`,
		14: `["twice a day", "half-day"]`,
		28: `["every six hours", "six-hrs"]`,
		56: `["every three hours", "three-hrs"]`,
	}
)

type UserCal struct {
	Email           string
	Calendars       []string
	UpdateIntervals []int
	Upcoming        []string
	Frequency       string `datastore:"-"`
	CalendarsJSON   string `datastore:"-"`
}

func init() {
	// Handles "/" or any other route not matched.
	http.HandleFunc("/", indexOr404)
}

func loginRedirect(w http.ResponseWriter, c appengine.Context, r *http.Request) {
	url, err := user.LoginURL(c, r.URL.String())
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	w.Header().Set("Location", url)
	w.WriteHeader(http.StatusFound)
}

func getUserCal(c appengine.Context, u *user.User) (*UserCal, error) {
	key := datastore.NewKey(c, "UserCal", u.ID, 0, nil)
	userCal := &UserCal{}
	err := datastore.Get(c, key, userCal)
	if err != nil {
		c.Infof("Nothing found for user: %v", u)
		c.Infof("Got an error: %v", err)
		baseInterval := 0 // TODO: Add logic.
		userCal = &UserCal{
			Email:           u.Email,
			Calendars:       []string{},
			UpdateIntervals: []int{baseInterval},
			Upcoming:        []string{},
			Frequency:       `["once a week", "week"]`,
			CalendarsJSON:   "[]",
		}
		_, err = datastore.Put(c, key, &userCal)
	} else {
		c.Infof("User was found: %v", u)
		// TODO: Implement PropertyLoadSaver interface.
		if userCal.Calendars == nil {
			userCal.Calendars = []string{}
		}
		if userCal.UpdateIntervals == nil {
			userCal.UpdateIntervals = []int{}
		}
		if userCal.Upcoming == nil {
			userCal.Upcoming = []string{}
		}

		var b []byte
		b, err = json.Marshal(userCal.Calendars)
		userCal.CalendarsJSON = string(b[:])
		userCal.Frequency = frequencyResponses[len(userCal.UpdateIntervals)]
		c.Infof("ent.Frequency: %#v, ent.CalendarsJSON: %#v",
			userCal.Frequency, userCal.CalendarsJSON)
	}

	return userCal, err
}

func index(writer http.ResponseWriter, request *http.Request) {
	c := appengine.NewContext(request)
	u := user.Current(c)
	if u == nil {
		loginRedirect(writer, c, request)
		return
	}

	userCal, err := getUserCal(c, u)
	if err != nil {
		http.Error(writer, err.Error(), http.StatusInternalServerError)
	}

	err = templateIndex.Execute(writer, userCal)
	if err != nil {
		http.Error(writer, err.Error(), http.StatusInternalServerError)
	}
}

func missing(writer http.ResponseWriter, unusedReq *http.Request) {
	writer.WriteHeader(http.StatusNotFound)
	err := template404.Execute(writer, nil)
	if err != nil {
		http.Error(writer, err.Error(), http.StatusInternalServerError)
	}
}

func indexOr404(writer http.ResponseWriter, request *http.Request) {
	if request.URL.Path == "/" {
		index(writer, request)
	} else {
		missing(writer, request)
	}
}
<commit_msg>Moving UserCal display logic into custom rendexIndex function.<commit_after>package main

// https://cloud.google.com/appengine/docs/go/tools/remoteapi
import _ "appengine/remote_api"

import (
	"encoding/json"
	"html/template"
	"net/http"

	"appengine"
	"appengine/datastore"
	"appengine/user"
)

var (
	templateIndex = template.Must(template.ParseFiles(
		"templates/index.html",
	))
	template404 = template.Must(template.ParseFiles(
		"templates/404.html",
	))
	frequencyResponses = map[int]string{
		1:  `["once a week", "week"]`,
		4:  `["every two days", "two-day"]`,
		7:  `["once a day", "day"]`,
		14: `["twice a day", "half-day"]`,
		28: `["every six hours", "six-hrs"]`,
		56: `["every three hours", "three-hrs"]`,
	}
)

type UserCal struct {
	Email           string
	Calendars       []string
	UpdateIntervals []int
	Upcoming        []string
	Frequency       string `datastore:"-"`
	CalendarsJSON   string `datastore:"-"`
}

func init() {
	// Handles "/" or any other route not matched.
	http.HandleFunc("/", indexOr404)
}

func loginRedirect(w http.ResponseWriter, c appengine.Context, r *http.Request) {
	url, err := user.LoginURL(c, r.URL.String())
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
	w.Header().Set("Location", url)
	w.WriteHeader(http.StatusFound)
}

func getUserCal(c appengine.Context, u *user.User) (*UserCal, error) {
	key := datastore.NewKey(c, "UserCal", u.ID, 0, nil)
	userCal := &UserCal{}
	err := datastore.Get(c, key, userCal)
	if err != nil {
		c.Infof("Nothing found for user: %v", u)
		c.Infof("Got an error: %v", err)
		baseInterval := 0 // TODO: Add logic.
		userCal = &UserCal{
			Email:           u.Email,
			Calendars:       []string{},
			UpdateIntervals: []int{baseInterval},
			Upcoming:        []string{},
			Frequency:       `["once a week", "week"]`,
			CalendarsJSON:   "[]",
		}
		_, err = datastore.Put(c, key, &userCal)
	} else {
		c.Infof("User was found: %v", u)
		// TODO: Implement PropertyLoadSaver interface.
		if userCal.Calendars == nil {
			userCal.Calendars = []string{}
		}
		if userCal.UpdateIntervals == nil {
			userCal.UpdateIntervals = []int{}
		}
		if userCal.Upcoming == nil {
			userCal.Upcoming = []string{}
		}
	}

	return userCal, err
}

func renderIndex(writer http.ResponseWriter, userCal *UserCal) error {
	b, err := json.Marshal(userCal.Calendars)
	if err != nil {
		return err
	}

	userCal.CalendarsJSON = string(b[:])
	userCal.Frequency = frequencyResponses[len(userCal.UpdateIntervals)]

	return templateIndex.Execute(writer, userCal)

}

func index(writer http.ResponseWriter, request *http.Request) {
	c := appengine.NewContext(request)
	u := user.Current(c)
	if u == nil {
		loginRedirect(writer, c, request)
		return
	}

	userCal, err := getUserCal(c, u)
	if err != nil {
		http.Error(writer, err.Error(), http.StatusInternalServerError)
	}

	err = renderIndex(writer, userCal)
	if err != nil {
		http.Error(writer, err.Error(), http.StatusInternalServerError)
	}
}

func missing(writer http.ResponseWriter, unusedReq *http.Request) {
	writer.WriteHeader(http.StatusNotFound)
	err := template404.Execute(writer, nil)
	if err != nil {
		http.Error(writer, err.Error(), http.StatusInternalServerError)
	}
}

func indexOr404(writer http.ResponseWriter, request *http.Request) {
	if request.URL.Path == "/" {
		index(writer, request)
	} else {
		missing(writer, request)
	}
}
<|endoftext|>


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

package file

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/discovery/targetgroup"
	yaml_util "github.com/prometheus/prometheus/util/yaml"
	"gopkg.in/fsnotify.v1"
	"gopkg.in/yaml.v2"
)

var (
	patFileSDName = regexp.MustCompile(`^[^*]*(\*[^/]*)?\.(json|yml|yaml|JSON|YML|YAML)$`)

	// DefaultSDConfig is the default file SD configuration.
	DefaultSDConfig = SDConfig{
		RefreshInterval: model.Duration(5 * time.Minute),
	}
)

// SDConfig is the configuration for file based discovery.
type SDConfig struct {
	Files           []string       `yaml:"files"`
	RefreshInterval model.Duration `yaml:"refresh_interval,omitempty"`

	// Catches all undefined fields and must be empty after parsing.
	XXX map[string]interface{} `yaml:",inline"`
}

// UnmarshalYAML implements the yaml.Unmarshaler interface.
func (c *SDConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
	*c = DefaultSDConfig
	type plain SDConfig
	err := unmarshal((*plain)(c))
	if err != nil {
		return err
	}
	if err := yaml_util.CheckOverflow(c.XXX, "file_sd_config"); err != nil {
		return err
	}
	if len(c.Files) == 0 {
		return fmt.Errorf("file service discovery config must contain at least one path name")
	}
	for _, name := range c.Files {
		if !patFileSDName.MatchString(name) {
			return fmt.Errorf("path name %q is not valid for file discovery", name)
		}
	}
	return nil
}

const fileSDFilepathLabel = model.MetaLabelPrefix + "filepath"

// TimestampCollector is a Custom Collector for Timestamps of the files.
type TimestampCollector struct {
	Description *prometheus.Desc
	discoverers map[*Discovery]struct{}
	lock        sync.RWMutex
}

// Describe method sends the description to the channel.
func (t *TimestampCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- t.Description
}

// Collect creates constant metrics for each file with last modified time of the file.
func (t *TimestampCollector) Collect(ch chan<- prometheus.Metric) {
	// New map to dedup filenames.
	uniqueFiles := make(map[string]float64)
	t.lock.RLock()
	for fileSD := range t.discoverers {
		for filename, timestamp := range fileSD.timestamps {
			uniqueFiles[filename] = timestamp
		}
	}
	t.lock.RUnlock()
	for filename, timestamp := range uniqueFiles {
		ch <- prometheus.MustNewConstMetric(
			t.Description,
			prometheus.GaugeValue,
			timestamp,
			filename,
		)
	}
}

func (t *TimestampCollector) addDiscoverer(disc *Discovery) {
	t.lock.Lock()
	t.discoverers[disc] = struct{}{}
	t.lock.Unlock()
}

func (t *TimestampCollector) removeDiscoverer(disc *Discovery) {
	t.lock.Lock()
	delete(t.discoverers, disc)
	t.lock.Unlock()
}

// NewTimestampCollector creates a TimestampCollector.
func NewTimestampCollector() *TimestampCollector {
	return &TimestampCollector{
		Description: prometheus.NewDesc(
			"prometheus_sd_file_timestamp",
			"Timestamp (mtime) of files read by FileSD. Timestamp is set at read time.",
			[]string{"filename"},
			nil,
		),
		discoverers: make(map[*Discovery]struct{}),
	}
}

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

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

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

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

// NewDiscovery returns a new file discovery for the given paths.
func NewDiscovery(conf *SDConfig, logger log.Logger) *Discovery {
	if logger == nil {
		logger = log.NewNopLogger()
	}

	disc := &Discovery{
		paths:      conf.Files,
		interval:   time.Duration(conf.RefreshInterval),
		timestamps: make(map[string]float64),
		logger:     logger,
	}
	fileSDTimeStamp.addDiscoverer(disc)
	return disc
}

// listFiles returns a list of all files that match the configured patterns.
func (d *Discovery) listFiles() []string {
	var paths []string
	for _, p := range d.paths {
		files, err := filepath.Glob(p)
		if err != nil {
			level.Error(d.logger).Log("msg", "Error expanding glob", "glob", p, "err", err)
			continue
		}
		paths = append(paths, files...)
	}
	return paths
}

// watchFiles sets watches on all full paths or directories that were configured for
// this file discovery.
func (d *Discovery) watchFiles() {
	if d.watcher == nil {
		panic("no watcher configured")
	}
	for _, p := range d.paths {
		if idx := strings.LastIndex(p, "/"); idx > -1 {
			p = p[:idx]
		} else {
			p = "./"
		}
		if err := d.watcher.Add(p); err != nil {
			level.Error(d.logger).Log("msg", "Error adding file watch", "path", p, "err", err)
		}
	}
}

// Run implements the Discoverer interface.
func (d *Discovery) Run(ctx context.Context, ch chan<- []*targetgroup.Group) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		level.Error(d.logger).Log("msg", "Error adding file watcher", "err", err)
		return
	}
	d.watcher = watcher
	defer d.stop()

	d.refresh(ctx, ch)

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

	for {
		select {
		case <-ctx.Done():
			return

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

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

		case err := <-d.watcher.Errors:
			if err != nil {
				level.Error(d.logger).Log("msg", "Error watching file", "err", err)
			}
		}
	}
}

func (d *Discovery) writeTimestamp(filename string, timestamp float64) {
	d.lock.Lock()
	d.timestamps[filename] = timestamp
	d.lock.Unlock()
}

func (d *Discovery) deleteTimestamp(filename string) {
	d.lock.Lock()
	delete(d.timestamps, filename)
	d.lock.Unlock()
}

// stop shuts down the file watcher.
func (d *Discovery) stop() {
	level.Debug(d.logger).Log("msg", "Stopping file discovery...", "paths", d.paths)

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

	fileSDTimeStamp.removeDiscoverer(d)

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

	level.Debug(d.logger).Log("File discovery stopped", "paths", d.paths)
}

// refresh reads all files matching the discovery's patterns and sends the respective
// updated target groups through the channel.
func (d *Discovery) refresh(ctx context.Context, ch chan<- []*targetgroup.Group) {
	t0 := time.Now()
	defer func() {
		fileSDScanDuration.Observe(time.Since(t0).Seconds())
	}()
	ref := map[string]int{}
	for _, p := range d.listFiles() {
		tgroups, err := d.readFile(p)
		if err != nil {
			fileSDReadErrorsCount.Inc()

			level.Error(d.logger).Log("msg", "Error reading file", "path", p, "err", err)
			// Prevent deletion down below.
			ref[p] = d.lastRefresh[p]
			continue
		}
		select {
		case ch <- tgroups:
		case <-ctx.Done():
			return
		}

		ref[p] = len(tgroups)
	}
	// Send empty updates for sources that disappeared.
	for f, n := range d.lastRefresh {
		m, ok := ref[f]
		if !ok || n > m {
			level.Debug(d.logger).Log("msg", "file_sd refresh found file that should be removed", "file", f)
			d.deleteTimestamp(f)
			for i := m; i < n; i++ {
				select {
				case ch <- []*targetgroup.Group{{Source: fileSource(f, i)}}:
				case <-ctx.Done():
					return
				}
			}
		}
	}
	d.lastRefresh = ref

	d.watchFiles()
}

// readFile reads a JSON or YAML list of targets groups from the file, depending on its
// file extension. It returns full configuration target groups.
func (d *Discovery) readFile(filename string) ([]*targetgroup.Group, error) {
	fd, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer fd.Close()

	content, err := ioutil.ReadAll(fd)
	if err != nil {
		return nil, err
	}

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

	var targetGroups []*targetgroup.Group

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

	for i, tg := range targetGroups {
		if tg == nil {
			err = errors.New("nil target group item found")
			return nil, err
		}

		tg.Source = fileSource(filename, i)
		if tg.Labels == nil {
			tg.Labels = model.LabelSet{}
		}
		tg.Labels[fileSDFilepathLabel] = model.LabelValue(filename)
	}

	d.writeTimestamp(filename, float64(info.ModTime().Unix()))

	return targetGroups, nil
}

// fileSource returns a source ID for the i-th target group in the file.
func fileSource(filename string, i int) string {
	return fmt.Sprintf("%s:%d", filename, i)
}
<commit_msg>discovery: Rename file SD mtime metric (#3723)<commit_after>// Copyright 2015 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package file

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/go-kit/kit/log"
	"github.com/go-kit/kit/log/level"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/discovery/targetgroup"
	yaml_util "github.com/prometheus/prometheus/util/yaml"
	"gopkg.in/fsnotify.v1"
	"gopkg.in/yaml.v2"
)

var (
	patFileSDName = regexp.MustCompile(`^[^*]*(\*[^/]*)?\.(json|yml|yaml|JSON|YML|YAML)$`)

	// DefaultSDConfig is the default file SD configuration.
	DefaultSDConfig = SDConfig{
		RefreshInterval: model.Duration(5 * time.Minute),
	}
)

// SDConfig is the configuration for file based discovery.
type SDConfig struct {
	Files           []string       `yaml:"files"`
	RefreshInterval model.Duration `yaml:"refresh_interval,omitempty"`

	// Catches all undefined fields and must be empty after parsing.
	XXX map[string]interface{} `yaml:",inline"`
}

// UnmarshalYAML implements the yaml.Unmarshaler interface.
func (c *SDConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
	*c = DefaultSDConfig
	type plain SDConfig
	err := unmarshal((*plain)(c))
	if err != nil {
		return err
	}
	if err := yaml_util.CheckOverflow(c.XXX, "file_sd_config"); err != nil {
		return err
	}
	if len(c.Files) == 0 {
		return fmt.Errorf("file service discovery config must contain at least one path name")
	}
	for _, name := range c.Files {
		if !patFileSDName.MatchString(name) {
			return fmt.Errorf("path name %q is not valid for file discovery", name)
		}
	}
	return nil
}

const fileSDFilepathLabel = model.MetaLabelPrefix + "filepath"

// TimestampCollector is a Custom Collector for Timestamps of the files.
type TimestampCollector struct {
	Description *prometheus.Desc
	discoverers map[*Discovery]struct{}
	lock        sync.RWMutex
}

// Describe method sends the description to the channel.
func (t *TimestampCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- t.Description
}

// Collect creates constant metrics for each file with last modified time of the file.
func (t *TimestampCollector) Collect(ch chan<- prometheus.Metric) {
	// New map to dedup filenames.
	uniqueFiles := make(map[string]float64)
	t.lock.RLock()
	for fileSD := range t.discoverers {
		for filename, timestamp := range fileSD.timestamps {
			uniqueFiles[filename] = timestamp
		}
	}
	t.lock.RUnlock()
	for filename, timestamp := range uniqueFiles {
		ch <- prometheus.MustNewConstMetric(
			t.Description,
			prometheus.GaugeValue,
			timestamp,
			filename,
		)
	}
}

func (t *TimestampCollector) addDiscoverer(disc *Discovery) {
	t.lock.Lock()
	t.discoverers[disc] = struct{}{}
	t.lock.Unlock()
}

func (t *TimestampCollector) removeDiscoverer(disc *Discovery) {
	t.lock.Lock()
	delete(t.discoverers, disc)
	t.lock.Unlock()
}

// NewTimestampCollector creates a TimestampCollector.
func NewTimestampCollector() *TimestampCollector {
	return &TimestampCollector{
		Description: prometheus.NewDesc(
			"prometheus_sd_file_mtime_seconds",
			"Timestamp (mtime) of files read by FileSD. Timestamp is set at read time.",
			[]string{"filename"},
			nil,
		),
		discoverers: make(map[*Discovery]struct{}),
	}
}

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

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

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

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

// NewDiscovery returns a new file discovery for the given paths.
func NewDiscovery(conf *SDConfig, logger log.Logger) *Discovery {
	if logger == nil {
		logger = log.NewNopLogger()
	}

	disc := &Discovery{
		paths:      conf.Files,
		interval:   time.Duration(conf.RefreshInterval),
		timestamps: make(map[string]float64),
		logger:     logger,
	}
	fileSDTimeStamp.addDiscoverer(disc)
	return disc
}

// listFiles returns a list of all files that match the configured patterns.
func (d *Discovery) listFiles() []string {
	var paths []string
	for _, p := range d.paths {
		files, err := filepath.Glob(p)
		if err != nil {
			level.Error(d.logger).Log("msg", "Error expanding glob", "glob", p, "err", err)
			continue
		}
		paths = append(paths, files...)
	}
	return paths
}

// watchFiles sets watches on all full paths or directories that were configured for
// this file discovery.
func (d *Discovery) watchFiles() {
	if d.watcher == nil {
		panic("no watcher configured")
	}
	for _, p := range d.paths {
		if idx := strings.LastIndex(p, "/"); idx > -1 {
			p = p[:idx]
		} else {
			p = "./"
		}
		if err := d.watcher.Add(p); err != nil {
			level.Error(d.logger).Log("msg", "Error adding file watch", "path", p, "err", err)
		}
	}
}

// Run implements the Discoverer interface.
func (d *Discovery) Run(ctx context.Context, ch chan<- []*targetgroup.Group) {
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		level.Error(d.logger).Log("msg", "Error adding file watcher", "err", err)
		return
	}
	d.watcher = watcher
	defer d.stop()

	d.refresh(ctx, ch)

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

	for {
		select {
		case <-ctx.Done():
			return

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

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

		case err := <-d.watcher.Errors:
			if err != nil {
				level.Error(d.logger).Log("msg", "Error watching file", "err", err)
			}
		}
	}
}

func (d *Discovery) writeTimestamp(filename string, timestamp float64) {
	d.lock.Lock()
	d.timestamps[filename] = timestamp
	d.lock.Unlock()
}

func (d *Discovery) deleteTimestamp(filename string) {
	d.lock.Lock()
	delete(d.timestamps, filename)
	d.lock.Unlock()
}

// stop shuts down the file watcher.
func (d *Discovery) stop() {
	level.Debug(d.logger).Log("msg", "Stopping file discovery...", "paths", d.paths)

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

	fileSDTimeStamp.removeDiscoverer(d)

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

	level.Debug(d.logger).Log("File discovery stopped", "paths", d.paths)
}

// refresh reads all files matching the discovery's patterns and sends the respective
// updated target groups through the channel.
func (d *Discovery) refresh(ctx context.Context, ch chan<- []*targetgroup.Group) {
	t0 := time.Now()
	defer func() {
		fileSDScanDuration.Observe(time.Since(t0).Seconds())
	}()
	ref := map[string]int{}
	for _, p := range d.listFiles() {
		tgroups, err := d.readFile(p)
		if err != nil {
			fileSDReadErrorsCount.Inc()

			level.Error(d.logger).Log("msg", "Error reading file", "path", p, "err", err)
			// Prevent deletion down below.
			ref[p] = d.lastRefresh[p]
			continue
		}
		select {
		case ch <- tgroups:
		case <-ctx.Done():
			return
		}

		ref[p] = len(tgroups)
	}
	// Send empty updates for sources that disappeared.
	for f, n := range d.lastRefresh {
		m, ok := ref[f]
		if !ok || n > m {
			level.Debug(d.logger).Log("msg", "file_sd refresh found file that should be removed", "file", f)
			d.deleteTimestamp(f)
			for i := m; i < n; i++ {
				select {
				case ch <- []*targetgroup.Group{{Source: fileSource(f, i)}}:
				case <-ctx.Done():
					return
				}
			}
		}
	}
	d.lastRefresh = ref

	d.watchFiles()
}

// readFile reads a JSON or YAML list of targets groups from the file, depending on its
// file extension. It returns full configuration target groups.
func (d *Discovery) readFile(filename string) ([]*targetgroup.Group, error) {
	fd, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer fd.Close()

	content, err := ioutil.ReadAll(fd)
	if err != nil {
		return nil, err
	}

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

	var targetGroups []*targetgroup.Group

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

	for i, tg := range targetGroups {
		if tg == nil {
			err = errors.New("nil target group item found")
			return nil, err
		}

		tg.Source = fileSource(filename, i)
		if tg.Labels == nil {
			tg.Labels = model.LabelSet{}
		}
		tg.Labels[fileSDFilepathLabel] = model.LabelValue(filename)
	}

	d.writeTimestamp(filename, float64(info.ModTime().Unix()))

	return targetGroups, nil
}

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