<commit_before>package som

import (
	"encoding/xml"
	"fmt"
	"io"
	"math"

	"github.com/gonum/matrix/mat64"
)

type rowWithDist struct {
	Row  int
	Dist float64
}

type h1 struct {
	XMLName xml.Name `xml:"h1"`
	Title   string   `xml:",innerxml"`
}

type polygon struct {
	XMLName xml.Name `xml:"polygon"`
	Points  []byte   `xml:"points,attr"`
	Style   string   `xml:"style,attr"`
}

type svgElement struct {
	XMLName  xml.Name `xml:"svg"`
	Width    float64  `xml:"width,attr"`
	Height   float64  `xml:"height,attr"`
	Polygons []interface{}
}

type textElement struct {
	XMLName xml.Name `xml:"text"`
	X       float64  `xml:"x,attr"`
	Y       float64  `xml:"y,attr"`
	Text    string   `xml:",innerxml"`
}

var colors = [][]int{{255, 0, 0}, {0, 255, 0}, {0, 0, 255}, {255, 255, 0}, {255, 0, 255}, {0, 255, 255}}

// Creates an SVG representation of the U-Matrix of the given codebook.
// codebook - the codebook we're displaying the U-Matrix for
// dims - the dimensions of the grid
// uShape - the shape of the grid
// title - the title of the output SVG
// writer - the io.Writter to write the output SVG to.
// clusters - if the clusters are known (i.e. these are test data) they can be displayed providing the information in this map.
// The map is: codebook vector row -> cluster number.  When clusters are not known (i.e. running with real data), just provide an empty map.
func UMatrixSVG(codebook *mat64.Dense, dims []int, uShape string, title string, writer io.Writer, clusters map[int]int) error {
	xmlEncoder := xml.NewEncoder(writer)
	// array to hold the xml elements
	elems := []interface{}{h1{Title: title}}

	rows, _ := codebook.Dims()
	distMat, err := DistanceMx("euclidean", codebook)
	if err != nil {
		return err
	}
	coords, err := GridCoords(uShape, dims)
	if err != nil {
		return err
	}
	coordsDistMat, err := DistanceMx("euclidean", coords)
	if err != nil {
		return err
	}

	umatrix := make([]float64, rows)
	maxDistance := -math.MaxFloat64
	minDistance := math.MaxFloat64
	for row := 0; row < rows; row++ {
		avgDistance := 0.0
		// this is a rough approximation of the notion of neighbor grid coords
		allRowsInRadius := allRowsInRadius(row, math.Sqrt2*1.01, coordsDistMat)
		for _, rwd := range allRowsInRadius {
			if rwd.Dist > 0.0 {
				avgDistance += distMat.At(row, rwd.Row)
			}
		}
		avgDistance /= float64(len(allRowsInRadius) - 1)
		umatrix[row] = avgDistance
		if avgDistance > maxDistance {
			maxDistance = avgDistance
		}
		if avgDistance < maxDistance {
			minDistance = avgDistance
		}
	}

	// function to scale the coord grid to something visible
	const MUL = 50.0
	const OFF = 10.0
	scale := func(x float64) float64 { return MUL*x + OFF }

	svgElem := svgElement{
		Width:    float64(dims[1])*MUL + 2*OFF,
		Height:   float64(dims[0])*MUL + 2*OFF,
		Polygons: make([]interface{}, rows*2),
	}
	for row := 0; row < rows; row++ {
		coord := coords.RowView(row)
		var colorMask []int
		clusterId, clusterFound := clusters[row]
		// if no cluster information, just use shades of gray
		if !clusterFound || clusterId == -1 {
			colorMask = []int{255, 255, 255}
		} else {
			colorMask = colors[clusters[row]%len(colors)]
		}
		colorMul := 1.0 - (umatrix[row]-minDistance)/(maxDistance-minDistance)
		r := int(colorMul * float64(colorMask[0]))
		g := int(colorMul * float64(colorMask[1]))
		b := int(colorMul * float64(colorMask[2]))
		polygonCoords := ""
		x := scale(coord.At(0, 0))
		y := scale(coord.At(1, 0))
		xOffset := 0.5 * MUL
		yOffset := 0.5 * MUL
		// hexagon has a different yOffset
		if uShape == "hexagon" {
			yOffset = math.Sqrt(0.75) / 2.0 * MUL
		}
		// draw a box around the current coord
		polygonCoords += fmt.Sprintf("%f,%f ", x+xOffset, y+yOffset)
		polygonCoords += fmt.Sprintf("%f,%f ", x+xOffset, y-yOffset)
		polygonCoords += fmt.Sprintf("%f,%f ", x-xOffset, y-yOffset)
		polygonCoords += fmt.Sprintf("%f,%f ", x-xOffset, y+yOffset)
		polygonCoords += fmt.Sprintf("%f,%f ", x+xOffset, y+yOffset)

		svgElem.Polygons[row*2] = polygon{
			Points: []byte(polygonCoords),
			Style:  fmt.Sprintf("fill:rgb(%d,%d,%d);stroke:black;stroke-width:1", r, g, b),
		}

		// print cluster number
		if clusterFound {
			svgElem.Polygons[row*2+1] = textElement{
				X:    x - 0.25*MUL,
				Y:    y + 0.25*MUL,
				Text: fmt.Sprintf("%d", clusters[row]),
			}
		}
	}

	elems = append(elems, svgElem)

	xmlEncoder.Encode(elems)
	xmlEncoder.Flush()

	return nil
}

func allRowsInRadius(selectedRow int, radius float64, distMatrix *mat64.Dense) []rowWithDist {
	rowsInRadius := []rowWithDist{}
	for i, dist := range distMatrix.RowView(selectedRow).RawVector().Data {
		if dist < radius {
			rowsInRadius = append(rowsInRadius, rowWithDist{Row: i, Dist: dist})
		}
	}
	return rowsInRadius
}
<commit_msg>improved readability<commit_after>package som

import (
	"encoding/xml"
	"fmt"
	"io"
	"math"

	"github.com/gonum/matrix/mat64"
)

type rowWithDist struct {
	Row  int
	Dist float64
}

type h1 struct {
	XMLName xml.Name `xml:"h1"`
	Title   string   `xml:",innerxml"`
}

type polygon struct {
	XMLName xml.Name `xml:"polygon"`
	Points  []byte   `xml:"points,attr"`
	Style   string   `xml:"style,attr"`
}

type svgElement struct {
	XMLName  xml.Name `xml:"svg"`
	Width    float64  `xml:"width,attr"`
	Height   float64  `xml:"height,attr"`
	Polygons []interface{}
}

type textElement struct {
	XMLName xml.Name `xml:"text"`
	X       float64  `xml:"x,attr"`
	Y       float64  `xml:"y,attr"`
	Text    string   `xml:",innerxml"`
}

var colors = [][]int{{255, 0, 0}, {0, 255, 0}, {0, 0, 255}, {255, 255, 0}, {255, 0, 255}, {0, 255, 255}}

// Creates an SVG representation of the U-Matrix of the given codebook.
// codebook - the codebook we're displaying the U-Matrix for
// dims - the dimensions of the grid
// uShape - the shape of the grid
// title - the title of the output SVG
// writer - the io.Writter to write the output SVG to.
// clusters - if the clusters are known (i.e. these are test data) they can be displayed providing the information in this map.
// The map is: codebook vector row -> cluster number.  When clusters are not known (i.e. running with real data), just provide an empty map.
func UMatrixSVG(codebook *mat64.Dense, dims []int, uShape, title string, writer io.Writer, clusters map[int]int) error {
	xmlEncoder := xml.NewEncoder(writer)
	// array to hold the xml elements
	elems := []interface{}{h1{Title: title}}

	rows, _ := codebook.Dims()
	distMat, err := DistanceMx("euclidean", codebook)
	if err != nil {
		return err
	}
	coords, err := GridCoords(uShape, dims)
	if err != nil {
		return err
	}
	coordsDistMat, err := DistanceMx("euclidean", coords)
	if err != nil {
		return err
	}

	umatrix := make([]float64, rows)
	maxDistance := -math.MaxFloat64
	minDistance := math.MaxFloat64
	for row := 0; row < rows; row++ {
		avgDistance := 0.0
		// this is a rough approximation of the notion of neighbor grid coords
		allRowsInRadius := allRowsInRadius(row, math.Sqrt2*1.01, coordsDistMat)
		for _, rwd := range allRowsInRadius {
			if rwd.Dist > 0.0 {
				avgDistance += distMat.At(row, rwd.Row)
			}
		}
		avgDistance /= float64(len(allRowsInRadius) - 1)
		umatrix[row] = avgDistance
		if avgDistance > maxDistance {
			maxDistance = avgDistance
		}
		if avgDistance < maxDistance {
			minDistance = avgDistance
		}
	}

	// function to scale the coord grid to something visible
	const MUL = 50.0
	const OFF = 10.0
	scale := func(x float64) float64 { return MUL*x + OFF }

	svgElem := svgElement{
		Width:    float64(dims[1])*MUL + 2*OFF,
		Height:   float64(dims[0])*MUL + 2*OFF,
		Polygons: make([]interface{}, rows*2),
	}
	for row := 0; row < rows; row++ {
		coord := coords.RowView(row)
		var colorMask []int
		clusterId, clusterFound := clusters[row]
		// if no cluster information, just use shades of gray
		if !clusterFound || clusterId == -1 {
			colorMask = []int{255, 255, 255}
		} else {
			colorMask = colors[clusters[row]%len(colors)]
		}
		colorMul := 1.0 - (umatrix[row]-minDistance)/(maxDistance-minDistance)
		r := int(colorMul * float64(colorMask[0]))
		g := int(colorMul * float64(colorMask[1]))
		b := int(colorMul * float64(colorMask[2]))
		polygonCoords := ""
		x := scale(coord.At(0, 0))
		y := scale(coord.At(1, 0))
		xOffset := 0.5 * MUL
		yOffset := 0.5 * MUL
		// hexagon has a different yOffset
		if uShape == "hexagon" {
			yOffset = math.Sqrt(0.75) / 2.0 * MUL
		}
		// draw a box around the current coord
		polygonCoords += fmt.Sprintf("%f,%f ", x+xOffset, y+yOffset)
		polygonCoords += fmt.Sprintf("%f,%f ", x+xOffset, y-yOffset)
		polygonCoords += fmt.Sprintf("%f,%f ", x-xOffset, y-yOffset)
		polygonCoords += fmt.Sprintf("%f,%f ", x-xOffset, y+yOffset)
		polygonCoords += fmt.Sprintf("%f,%f ", x+xOffset, y+yOffset)

		svgElem.Polygons[row*2] = polygon{
			Points: []byte(polygonCoords),
			Style:  fmt.Sprintf("fill:rgb(%d,%d,%d);stroke:black;stroke-width:1", r, g, b),
		}

		// print cluster number
		if clusterFound {
			svgElem.Polygons[row*2+1] = textElement{
				X:    x - 0.25*MUL,
				Y:    y + 0.25*MUL,
				Text: fmt.Sprintf("%d", clusters[row]),
			}
		}
	}

	elems = append(elems, svgElem)

	xmlEncoder.Encode(elems)
	xmlEncoder.Flush()

	return nil
}

func allRowsInRadius(selectedRow int, radius float64, distMatrix *mat64.Dense) []rowWithDist {
	rowsInRadius := []rowWithDist{}
	for i, dist := range distMatrix.RowView(selectedRow).RawVector().Data {
		if dist < radius {
			rowsInRadius = append(rowsInRadius, rowWithDist{Row: i, Dist: dist})
		}
	}
	return rowsInRadius
}
<|endoftext|>


<commit_before>// +build !windows

// Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 statemanager

import (
	"io/ioutil"
	"os"
	"path/filepath"
)

/*
On Linux, the basic approach for attempting to ensure that the state file is
written out correctly relies on behaviors of Linux and the ext* family of
filesystems.

On each save, the agent creates a new temporary file where it
writes out the json object.  Once the file is written, it gets renamed to the
well-known name of the state file.  Under the assumption of Linux + ext*, this
is an atomic operation; rename is changing the hard link of the well-known file
to point to the inode of the temporary file.  The original file inode now has no
links, and is considered free space once the opened file handles to that inode
are closed.

On each load, the agent opens a well-known file name for the state file and
reads it.
*/

func newPlatformDependencies() platformDependencies {
	return nil
}

func (manager *basicStateManager) readFile() ([]byte, error) {
	// Note that even if Save overwrites the file we're looking at here, we
	// still hold the old inode and should read the old data so no locking is
	// needed (given Linux and the ext* family of fs at least).
	file, err := os.Open(filepath.Join(manager.statePath, ecsDataFile))
	if err != nil {
		if os.IsNotExist(err) {
			// Happens every first run; not a real error
			return nil, nil
		}
		return nil, err
	}
	return ioutil.ReadAll(file)
}

func (manager *basicStateManager) writeFile(data []byte) error {
	// Make our temp-file on the same volume as our data-file to ensure we can
	// actually move it atomically; cross-device renaming will error out.
	tmpfile, err := ioutil.TempFile(manager.statePath, "tmp_ecs_agent_data")
	if err != nil {
		log.Error("Error saving state; could not create temp file to save state", "err", err)
		return err
	}
	_, err = tmpfile.Write(data)
	if err != nil {
		log.Error("Error saving state; could not write to temp file to save state", "err", err)
		return err
	}
	err = os.Rename(tmpfile.Name(), filepath.Join(manager.statePath, ecsDataFile))
	if err != nil {
		log.Error("Error saving state; could not move to data file", "err", err)
	}
	return err
}
<commit_msg>Revert "Revert "sync temp state file after write""<commit_after>// +build !windows

// Copyright 2014-2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file 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 statemanager

import (
	"io/ioutil"
	"os"
	"path/filepath"

	"github.com/cihub/seelog"
)

/*
On Linux, the basic approach for attempting to ensure that the state file is
written out correctly relies on behaviors of Linux and the ext* family of
filesystems.

On each save, the agent creates a new temporary file where it
writes out the json object.  Once the file is written, it gets renamed to the
well-known name of the state file.  Under the assumption of Linux + ext*, this
is an atomic operation; rename is changing the hard link of the well-known file
to point to the inode of the temporary file.  The original file inode now has no
links, and is considered free space once the opened file handles to that inode
are closed.

On each load, the agent opens a well-known file name for the state file and
reads it.
*/

func newPlatformDependencies() platformDependencies {
	return nil
}

func (manager *basicStateManager) readFile() ([]byte, error) {
	// Note that even if Save overwrites the file we're looking at here, we
	// still hold the old inode and should read the old data so no locking is
	// needed (given Linux and the ext* family of fs at least).
	file, err := os.Open(filepath.Join(manager.statePath, ecsDataFile))
	if err != nil {
		if os.IsNotExist(err) {
			// Happens every first run; not a real error
			return nil, nil
		}
		return nil, err
	}
	return ioutil.ReadAll(file)
}

func (manager *basicStateManager) writeFile(data []byte) error {
	// Make our temp-file on the same volume as our data-file to ensure we can
	// actually move it atomically; cross-device renaming will error out.
	tmpfile, err := ioutil.TempFile(manager.statePath, "tmp_ecs_agent_data")
	if err != nil {
		seelog.Errorf("Error saving state; could not create temp file to save state, err: %v", err)
		return err
	}
	_, err = tmpfile.Write(data)
	if err != nil {
		seelog.Errorf("Error saving state; could not write to temp file to save state, err: %v", err)
		return err
	}

	// flush temp state file to disk
	err = tmpfile.Sync()
	if err != nil {
		seelog.Errorf("Error flusing state file, err: %v", err)
		return err
	}

	err = os.Rename(tmpfile.Name(), filepath.Join(manager.statePath, ecsDataFile))
	if err != nil {
		seelog.Errorf("Error saving state; could not move to data file, err: %v", err)
	}
	return err
}
<|endoftext|>


<commit_before>package tcp // import "github.com/influxdata/influxdb/tcp"

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"sync"
	"time"
)

const (
	// DefaultTimeout is the default length of time to wait for first byte.
	DefaultTimeout = 30 * time.Second
)

// Mux multiplexes a network connection.
type Mux struct {
	mu sync.RWMutex
	ln net.Listener
	m  map[byte]*listener

	defaultListener *listener

	wg sync.WaitGroup

	// The amount of time to wait for the first header byte.
	Timeout time.Duration

	// Out-of-band error logger
	Logger *log.Logger
}

type replayConn struct {
	net.Conn
	firstByte     byte
	readFirstbyte bool
}

func (rc *replayConn) Read(b []byte) (int, error) {
	if rc.readFirstbyte {
		return rc.Conn.Read(b)
	}

	if len(b) == 0 {
		return 0, nil
	}

	b[0] = rc.firstByte
	rc.readFirstbyte = true
	return 1, nil
}

// NewMux returns a new instance of Mux for ln.
func NewMux() *Mux {
	return &Mux{
		m:       make(map[byte]*listener),
		Timeout: DefaultTimeout,
		Logger:  log.New(os.Stderr, "[tcp] ", log.LstdFlags),
	}
}

// Serve handles connections from ln and multiplexes then across registered listener.
func (mux *Mux) Serve(ln net.Listener) error {
	mux.mu.Lock()
	mux.ln = ln
	mux.mu.Unlock()
	for {
		// Wait for the next connection.
		// If it returns a temporary error then simply retry.
		// If it returns any other error then exit immediately.
		conn, err := ln.Accept()
		if err, ok := err.(interface {
			Temporary() bool
		}); ok && err.Temporary() {
			continue
		}
		if err != nil {
			// Wait for all connections to be demux
			mux.wg.Wait()
			for _, ln := range mux.m {
				close(ln.c)
			}

			if mux.defaultListener != nil {
				close(mux.defaultListener.c)
			}

			return err
		}

		// Demux in a goroutine to
		mux.wg.Add(1)
		go mux.handleConn(conn)
	}
}

func (mux *Mux) handleConn(conn net.Conn) {
	defer mux.wg.Done()
	// Set a read deadline so connections with no data don't timeout.
	if err := conn.SetReadDeadline(time.Now().Add(mux.Timeout)); err != nil {
		conn.Close()
		mux.Logger.Printf("tcp.Mux: cannot set read deadline: %s", err)
		return
	}

	// Read first byte from connection to determine handler.
	var typ [1]byte
	if _, err := io.ReadFull(conn, typ[:]); err != nil {
		conn.Close()
		mux.Logger.Printf("tcp.Mux: cannot read header byte: %s", err)
		return
	}

	// Reset read deadline and let the listener handle that.
	if err := conn.SetReadDeadline(time.Time{}); err != nil {
		conn.Close()
		mux.Logger.Printf("tcp.Mux: cannot reset set read deadline: %s", err)
		return
	}

	// Retrieve handler based on first byte.
	handler := mux.m[typ[0]]
	if handler == nil {
		if mux.defaultListener == nil {
			conn.Close()
			mux.Logger.Printf("tcp.Mux: handler not registered: %d. Connection from %s closed", typ[0], conn.RemoteAddr())
			return
		}

		conn = &replayConn{
			Conn:      conn,
			firstByte: typ[0],
		}
		handler = mux.defaultListener
	}

	// Send connection to handler.  The handler is responsible for closing the connection.
	timer := time.NewTimer(mux.Timeout)
	defer timer.Stop()

	select {
	case handler.c <- conn:
	case <-timer.C:
		conn.Close()
		mux.Logger.Printf("tcp.Mux: handler not ready: %d. Connection from %s closed", typ[0], conn.RemoteAddr())
		return
	}
}

// Listen returns a listener identified by header.
// Any connection accepted by mux is multiplexed based on the initial header byte.
func (mux *Mux) Listen(header byte) net.Listener {
	// Ensure two listeners are not created for the same header byte.
	if _, ok := mux.m[header]; ok {
		panic(fmt.Sprintf("listener already registered under header byte: %d", header))
	}

	// Create a new listener and assign it.
	ln := &listener{
		c:   make(chan net.Conn),
		mux: mux,
	}
	mux.m[header] = ln

	return ln
}

// DefaultListener() will return a net.Listener that will pass-through any
// connections with non-registered values for the first byte of the connection.
// The connections returned from this listener's Accept() method will replay the
// first byte of the connection as a short first Read().
//
// This can be used to pass to an HTTP server, so long as there are no conflicts
// with registsered listener bytes and the first character of the HTTP request:
// 71 ('G') for GET, etc.
func (mux *Mux) DefaultListener() net.Listener {
	if mux.defaultListener == nil {
		mux.defaultListener = &listener{
			c:   make(chan net.Conn),
			mux: mux,
		}
	}

	return mux.defaultListener
}

// listener is a receiver for connections received by Mux.
type listener struct {
	c   chan net.Conn
	mux *Mux
}

// Accept waits for and returns the next connection to the listener.
func (ln *listener) Accept() (c net.Conn, err error) {
	conn, ok := <-ln.c
	if !ok {
		return nil, errors.New("network connection closed")
	}
	return conn, nil
}

// Close is a no-op. The mux's listener should be closed instead.
func (ln *listener) Close() error { return nil }

// Addr returns the Addr of the listener
func (ln *listener) Addr() net.Addr {
	if ln.mux == nil {
		return nil
	}

	ln.mux.mu.RLock()
	defer ln.mux.mu.RUnlock()

	if ln.mux.ln == nil {
		return nil
	}

	return ln.mux.ln.Addr()
}

// Dial connects to a remote mux listener with a given header byte.
func Dial(network, address string, header byte) (net.Conn, error) {
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}

	if _, err := conn.Write([]byte{header}); err != nil {
		return nil, fmt.Errorf("write mux header: %s", err)
	}

	return conn, nil
}
<commit_msg>removed parentheses to make tcp golintable<commit_after>package tcp // import "github.com/influxdata/influxdb/tcp"

import (
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"sync"
	"time"
)

const (
	// DefaultTimeout is the default length of time to wait for first byte.
	DefaultTimeout = 30 * time.Second
)

// Mux multiplexes a network connection.
type Mux struct {
	mu sync.RWMutex
	ln net.Listener
	m  map[byte]*listener

	defaultListener *listener

	wg sync.WaitGroup

	// The amount of time to wait for the first header byte.
	Timeout time.Duration

	// Out-of-band error logger
	Logger *log.Logger
}

type replayConn struct {
	net.Conn
	firstByte     byte
	readFirstbyte bool
}

func (rc *replayConn) Read(b []byte) (int, error) {
	if rc.readFirstbyte {
		return rc.Conn.Read(b)
	}

	if len(b) == 0 {
		return 0, nil
	}

	b[0] = rc.firstByte
	rc.readFirstbyte = true
	return 1, nil
}

// NewMux returns a new instance of Mux for ln.
func NewMux() *Mux {
	return &Mux{
		m:       make(map[byte]*listener),
		Timeout: DefaultTimeout,
		Logger:  log.New(os.Stderr, "[tcp] ", log.LstdFlags),
	}
}

// Serve handles connections from ln and multiplexes then across registered listener.
func (mux *Mux) Serve(ln net.Listener) error {
	mux.mu.Lock()
	mux.ln = ln
	mux.mu.Unlock()
	for {
		// Wait for the next connection.
		// If it returns a temporary error then simply retry.
		// If it returns any other error then exit immediately.
		conn, err := ln.Accept()
		if err, ok := err.(interface {
			Temporary() bool
		}); ok && err.Temporary() {
			continue
		}
		if err != nil {
			// Wait for all connections to be demux
			mux.wg.Wait()
			for _, ln := range mux.m {
				close(ln.c)
			}

			if mux.defaultListener != nil {
				close(mux.defaultListener.c)
			}

			return err
		}

		// Demux in a goroutine to
		mux.wg.Add(1)
		go mux.handleConn(conn)
	}
}

func (mux *Mux) handleConn(conn net.Conn) {
	defer mux.wg.Done()
	// Set a read deadline so connections with no data don't timeout.
	if err := conn.SetReadDeadline(time.Now().Add(mux.Timeout)); err != nil {
		conn.Close()
		mux.Logger.Printf("tcp.Mux: cannot set read deadline: %s", err)
		return
	}

	// Read first byte from connection to determine handler.
	var typ [1]byte
	if _, err := io.ReadFull(conn, typ[:]); err != nil {
		conn.Close()
		mux.Logger.Printf("tcp.Mux: cannot read header byte: %s", err)
		return
	}

	// Reset read deadline and let the listener handle that.
	if err := conn.SetReadDeadline(time.Time{}); err != nil {
		conn.Close()
		mux.Logger.Printf("tcp.Mux: cannot reset set read deadline: %s", err)
		return
	}

	// Retrieve handler based on first byte.
	handler := mux.m[typ[0]]
	if handler == nil {
		if mux.defaultListener == nil {
			conn.Close()
			mux.Logger.Printf("tcp.Mux: handler not registered: %d. Connection from %s closed", typ[0], conn.RemoteAddr())
			return
		}

		conn = &replayConn{
			Conn:      conn,
			firstByte: typ[0],
		}
		handler = mux.defaultListener
	}

	// Send connection to handler.  The handler is responsible for closing the connection.
	timer := time.NewTimer(mux.Timeout)
	defer timer.Stop()

	select {
	case handler.c <- conn:
	case <-timer.C:
		conn.Close()
		mux.Logger.Printf("tcp.Mux: handler not ready: %d. Connection from %s closed", typ[0], conn.RemoteAddr())
		return
	}
}

// Listen returns a listener identified by header.
// Any connection accepted by mux is multiplexed based on the initial header byte.
func (mux *Mux) Listen(header byte) net.Listener {
	// Ensure two listeners are not created for the same header byte.
	if _, ok := mux.m[header]; ok {
		panic(fmt.Sprintf("listener already registered under header byte: %d", header))
	}

	// Create a new listener and assign it.
	ln := &listener{
		c:   make(chan net.Conn),
		mux: mux,
	}
	mux.m[header] = ln

	return ln
}

// DefaultListener will return a net.Listener that will pass-through any
// connections with non-registered values for the first byte of the connection.
// The connections returned from this listener's Accept() method will replay the
// first byte of the connection as a short first Read().
//
// This can be used to pass to an HTTP server, so long as there are no conflicts
// with registsered listener bytes and the first character of the HTTP request:
// 71 ('G') for GET, etc.
func (mux *Mux) DefaultListener() net.Listener {
	if mux.defaultListener == nil {
		mux.defaultListener = &listener{
			c:   make(chan net.Conn),
			mux: mux,
		}
	}

	return mux.defaultListener
}

// listener is a receiver for connections received by Mux.
type listener struct {
	c   chan net.Conn
	mux *Mux
}

// Accept waits for and returns the next connection to the listener.
func (ln *listener) Accept() (c net.Conn, err error) {
	conn, ok := <-ln.c
	if !ok {
		return nil, errors.New("network connection closed")
	}
	return conn, nil
}

// Close is a no-op. The mux's listener should be closed instead.
func (ln *listener) Close() error { return nil }

// Addr returns the Addr of the listener
func (ln *listener) Addr() net.Addr {
	if ln.mux == nil {
		return nil
	}

	ln.mux.mu.RLock()
	defer ln.mux.mu.RUnlock()

	if ln.mux.ln == nil {
		return nil
	}

	return ln.mux.ln.Addr()
}

// Dial connects to a remote mux listener with a given header byte.
func Dial(network, address string, header byte) (net.Conn, error) {
	conn, err := net.Dial(network, address)
	if err != nil {
		return nil, err
	}

	if _, err := conn.Write([]byte{header}); err != nil {
		return nil, fmt.Errorf("write mux header: %s", err)
	}

	return conn, nil
}
<|endoftext|>


<commit_before>// Copyright 2015 The Hugo Authors. All rights reserved.
//
// 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 source

import (
	"io"
	"path/filepath"
	"strings"

	"github.com/spf13/hugo/hugofs"

	"github.com/spf13/hugo/config"
	"github.com/spf13/hugo/helpers"
)

type SourceSpec struct {
	Cfg config.Provider
	Fs  *hugofs.Fs
}

func NewSourceSpec(cfg config.Provider, fs *hugofs.Fs) SourceSpec {
	return SourceSpec{Cfg: cfg, Fs: fs}
}

// File represents a source content file.
// All paths are relative from the source directory base
type File struct {
	relpath     string // Original relative path, e.g. section/foo.txt
	logicalName string // foo.txt
	baseName    string // `post` for `post.md`, also `post.en` for `post.en.md`
	Contents    io.Reader
	section     string // The first directory
	dir         string // The relative directory Path (minus file name)
	ext         string // Just the ext (eg txt)
	uniqueID    string // MD5 of the file's path

	translationBaseName string // `post` for `post.es.md` (if `Multilingual` is enabled.)
	lang                string // The language code if `Multilingual` is enabled
}

// UniqueID is the MD5 hash of the file's path and is for most practical applications,
// Hugo content files being one of them, considered to be unique.
func (f *File) UniqueID() string {
	return f.uniqueID
}

// String returns the file's content as a string.
func (f *File) String() string {
	return helpers.ReaderToString(f.Contents)
}

// Bytes returns the file's content as a byte slice.
func (f *File) Bytes() []byte {
	return helpers.ReaderToBytes(f.Contents)
}

// BaseFileName is a filename without extension.
func (f *File) BaseFileName() string {
	return f.baseName
}

// TranslationBaseName is a filename with no extension,
// not even the optional language extension part.
func (f *File) TranslationBaseName() string {
	return f.translationBaseName
}

// Lang for this page, if `Multilingual` is enabled on your site.
func (f *File) Lang() string {
	return f.lang
}

// Section is first directory below the content root.
func (f *File) Section() string {
	return f.section
}

// LogicalName is filename and extension of the file.
func (f *File) LogicalName() string {
	return f.logicalName
}

// SetDir sets the relative directory where this file lives.
// TODO(bep) Get rid of this.
func (f *File) SetDir(dir string) {
	f.dir = dir
}

// Dir gets the name of the directory that contains this file.
// The directory is relative to the content root.
func (f *File) Dir() string {
	return f.dir
}

// Extension gets the file extension, i.e "myblogpost.md" will return "md".
func (f *File) Extension() string {
	return f.ext
}

// Ext is an alias for Extension.
func (f *File) Ext() string {
	return f.Extension()
}

// Path gets the relative path including file name and extension.
// The directory is relative to the content root.
func (f *File) Path() string {
	return f.relpath
}

// NewFileWithContents creates a new File pointer with the given relative path and
// content. The language defaults to "en".
func (sp SourceSpec) NewFileWithContents(relpath string, content io.Reader) *File {
	file := sp.NewFile(relpath)
	file.Contents = content
	file.lang = "en"
	return file
}

// NewFile creates a new File pointer with the given relative path.
func (sp SourceSpec) NewFile(relpath string) *File {
	f := &File{
		relpath: relpath,
	}

	f.dir, f.logicalName = filepath.Split(f.relpath)
	f.ext = strings.TrimPrefix(filepath.Ext(f.LogicalName()), ".")
	f.baseName = helpers.Filename(f.LogicalName())

	lang := strings.TrimPrefix(filepath.Ext(f.baseName), ".")
	if _, ok := sp.Cfg.GetStringMap("languages")[lang]; lang == "" || !ok {
		f.lang = sp.Cfg.GetString("defaultContentLanguage")
		f.translationBaseName = f.baseName
	} else {
		f.lang = lang
		f.translationBaseName = helpers.Filename(f.baseName)
	}

	f.section = helpers.GuessSection(f.Dir())
	f.uniqueID = helpers.Md5String(f.Path())

	return f
}

// NewFileFromAbs creates a new File pointer with the given full file path path and
// content.
func (sp SourceSpec) NewFileFromAbs(base, fullpath string, content io.Reader) (f *File, err error) {
	var name string
	if name, err = helpers.GetRelativePath(fullpath, base); err != nil {
		return nil, err
	}

	return sp.NewFileWithContents(name, content), nil
}
<commit_msg>source: Cache language config<commit_after>// Copyright 2015 The Hugo Authors. All rights reserved.
//
// 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 source

import (
	"io"
	"path/filepath"
	"strings"

	"github.com/spf13/hugo/hugofs"

	"github.com/spf13/hugo/config"
	"github.com/spf13/hugo/helpers"
)

type SourceSpec struct {
	Cfg config.Provider
	Fs  *hugofs.Fs

	languages              map[string]interface{}
	defaultContentLanguage string
}

func NewSourceSpec(cfg config.Provider, fs *hugofs.Fs) SourceSpec {
	defaultLang := cfg.GetString("defaultContentLanguage")
	languages := cfg.GetStringMap("languages")
	return SourceSpec{Cfg: cfg, Fs: fs, languages: languages, defaultContentLanguage: defaultLang}
}

// File represents a source content file.
// All paths are relative from the source directory base
type File struct {
	relpath     string // Original relative path, e.g. section/foo.txt
	logicalName string // foo.txt
	baseName    string // `post` for `post.md`, also `post.en` for `post.en.md`
	Contents    io.Reader
	section     string // The first directory
	dir         string // The relative directory Path (minus file name)
	ext         string // Just the ext (eg txt)
	uniqueID    string // MD5 of the file's path

	translationBaseName string // `post` for `post.es.md` (if `Multilingual` is enabled.)
	lang                string // The language code if `Multilingual` is enabled
}

// UniqueID is the MD5 hash of the file's path and is for most practical applications,
// Hugo content files being one of them, considered to be unique.
func (f *File) UniqueID() string {
	return f.uniqueID
}

// String returns the file's content as a string.
func (f *File) String() string {
	return helpers.ReaderToString(f.Contents)
}

// Bytes returns the file's content as a byte slice.
func (f *File) Bytes() []byte {
	return helpers.ReaderToBytes(f.Contents)
}

// BaseFileName is a filename without extension.
func (f *File) BaseFileName() string {
	return f.baseName
}

// TranslationBaseName is a filename with no extension,
// not even the optional language extension part.
func (f *File) TranslationBaseName() string {
	return f.translationBaseName
}

// Lang for this page, if `Multilingual` is enabled on your site.
func (f *File) Lang() string {
	return f.lang
}

// Section is first directory below the content root.
func (f *File) Section() string {
	return f.section
}

// LogicalName is filename and extension of the file.
func (f *File) LogicalName() string {
	return f.logicalName
}

// SetDir sets the relative directory where this file lives.
// TODO(bep) Get rid of this.
func (f *File) SetDir(dir string) {
	f.dir = dir
}

// Dir gets the name of the directory that contains this file.
// The directory is relative to the content root.
func (f *File) Dir() string {
	return f.dir
}

// Extension gets the file extension, i.e "myblogpost.md" will return "md".
func (f *File) Extension() string {
	return f.ext
}

// Ext is an alias for Extension.
func (f *File) Ext() string {
	return f.Extension()
}

// Path gets the relative path including file name and extension.
// The directory is relative to the content root.
func (f *File) Path() string {
	return f.relpath
}

// NewFileWithContents creates a new File pointer with the given relative path and
// content. The language defaults to "en".
func (sp SourceSpec) NewFileWithContents(relpath string, content io.Reader) *File {
	file := sp.NewFile(relpath)
	file.Contents = content
	file.lang = "en"
	return file
}

// NewFile creates a new File pointer with the given relative path.
func (sp SourceSpec) NewFile(relpath string) *File {
	f := &File{
		relpath: relpath,
	}

	f.dir, f.logicalName = filepath.Split(f.relpath)
	f.ext = strings.TrimPrefix(filepath.Ext(f.LogicalName()), ".")
	f.baseName = helpers.Filename(f.LogicalName())

	lang := strings.TrimPrefix(filepath.Ext(f.baseName), ".")
	if _, ok := sp.languages[lang]; lang == "" || !ok {
		f.lang = sp.defaultContentLanguage
		f.translationBaseName = f.baseName
	} else {
		f.lang = lang
		f.translationBaseName = helpers.Filename(f.baseName)
	}

	f.section = helpers.GuessSection(f.Dir())
	f.uniqueID = helpers.Md5String(f.Path())

	return f
}

// NewFileFromAbs creates a new File pointer with the given full file path path and
// content.
func (sp SourceSpec) NewFileFromAbs(base, fullpath string, content io.Reader) (f *File, err error) {
	var name string
	if name, err = helpers.GetRelativePath(fullpath, base); err != nil {
		return nil, err
	}

	return sp.NewFileWithContents(name, content), nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"hash/crc64"
	"io"
	"math/rand"
	"os"
	"os/exec"
	"runtime"
	"sync"
	"testing"
	"time"
)

func TestSigpipe(t *testing.T) {
	r, w, err := os.Pipe()
	if err != nil {
		t.Fatal(err)
	}
	cmd := exec.Command("head", "-c64", "/dev/urandom")
	cmd.Stdout = w
	if err = cmd.Start(); err != nil {
		t.Fatal(err)
	}
	w.Close()
	go cmd.Wait()
	source := GetSource(fmt.Sprintf("/dev/fd/%d", r.Fd()), &Config{
		SourceBuffer: 5,
		FrameBytes:   16,
		CloseIdle:    false,
		Reopen:       true,
	})
	var frame = make(DataFrame, source.frameBytes)
	var nextFrame uint64
	for f := 0; f < 4; f++ {
		if _, err = source.Next(&nextFrame, &frame); err != nil {
			t.Error(err)
			break
		}
	}
	done := make(chan error, 1)
	go func() {
		_, err = source.Next(&nextFrame, &frame)
		done <- err
	}()
	select {
	case <-time.After(time.Second):
		t.Error("Should have given up within 1s of SIGCHLD")
	case err = <-done:
		if err == nil {
			t.Errorf("Should have error, got frame %v", frame)
		}
	}
	r.Close()
	source.Done()
	CloseAllSources()
}

func TestEmptySource(t *testing.T) {
	r, w, err := os.Pipe()
	if err != nil {
		t.Fatal(err)
	}
	closeNow := make(chan bool)
	go func() {
		<-closeNow
		w.Close()
	}()
	source := GetSource(fmt.Sprintf("/dev/fd/%d", r.Fd()), &Config{
		SourceBuffer: 5,
		FrameBytes:   16,
		CloseIdle:    false,
		Reopen:       false,
	})
	var frame = make(DataFrame, source.frameBytes)
	var nextFrame uint64
	done := make(chan error, 1)
	go func() {
		_, err := source.Next(&nextFrame, &frame)
		done <- err
	}()
	time.Sleep(time.Millisecond)
	select {
	case err := <-done:
		t.Errorf("Should still be waiting for input, got error %v", err)
	default:
	}
	closeNow <- true
	select {
	case <-time.After(100 * time.Millisecond):
		t.Error("Should have given up within 100ms")
	case err := <-done:
		if err == nil {
			t.Errorf("Should have error, got frame %v", frame)
		}
	}
	source.Done()
	r.Close()
	CloseAllSources()
}

func failUnless(t *testing.T, ms time.Duration, ok <-chan bool) {
	select {
	case <-time.After(ms * time.Millisecond):
		t.Fatalf("Timed out in %d ms", ms)
	case <-ok:
	}
}

func TestSentEqualsReceived(t *testing.T) {
	r, w, err := os.Pipe()
	if err != nil {
		t.Fatal(err)
	}
	sentData := []byte{}
	wantMore := make(chan int)
	go func() {
		for n := range wantMore {
			if n == 0 {
				break
			}
			moreData := make([]byte, n)
			for i := 0; i < n; i++ {
				moreData[i] = byte(rand.Int() & 0xff)
			}
			sentData = append(sentData, moreData...)
			if did, err := w.Write(moreData); did < len(moreData) || err != nil {
				t.Error("Short write: %d < %d, %s", did, len(moreData), err)
			}
			w.Sync()
		}
		w.Close()
	}()
	source := GetSource(fmt.Sprintf("/dev/fd/%d", r.Fd()), &Config{
		SourceBuffer: 5,
		FrameBytes:   16,
		CloseIdle:    true,
		Reopen:       false,
	})
	var rcvdData = []byte{}
	var frame = make(DataFrame, source.frameBytes)
	var nextFrame uint64
	ok := make(chan bool)
	wantMore <- 19 // Send one full frame (to make sure we start receiving at frame 0) and a bit extra
	for f := 100; f > 0; f-- {
		go func() {
			if _, err := source.Next(&nextFrame, &frame); err != nil {
				t.Fatal(err)
			}
			if len(frame) != 16 {
				t.Errorf("Wrong size frame, want 16 got %d", len(frame))
			}
			ok <- true
		}()
		failUnless(t, 100, ok)
		rcvdData = append(rcvdData, frame...)
		if f == 4 {
			wantMore <- 16 * 3 // Last three frames (we sent one frame before the loop)
		} else if f%4 == 0 {
			wantMore <- 16 * 4 // Next four frames
		}
	}
	close(wantMore)
	go func() {
		if _, err := source.Next(&nextFrame, &frame); err != io.EOF {
			t.Error("Should have got EOF")
		}
		ok <- true
	}()
	failUnless(t, 100, ok)
	wantData := sentData[0:1600]
	if 0 != bytes.Compare(wantData, rcvdData) {
		t.Errorf("want %d != rcvd %d", len(wantData), len(rcvdData))
	}
	source.Done()
	r.Close()
	CloseAllSources()
}

func TestHeader(t *testing.T) {
	headerSize := uint64(64)
	nClients := 5
	sources := make(chan *Source, nClients)
	for i := 0; i < nClients; i++ {
		sources <- GetSource("/dev/urandom", &Config{
			SourceBuffer: 5,
			FrameBytes:   65536,
			HeaderBytes:  headerSize,
			CloseIdle:    true,
		})
	}
	empty := make([]byte, headerSize)
	var h0 []byte
	for i := 0; i < nClients; i++ {
		h := make([]byte, headerSize)
		source := <-sources
		err := source.GetHeader(h)
		if err != nil {
			t.Error(err)
		} else if h0 == nil {
			h0 = h
		} else if bytes.Compare(h, h0) != 0 {
			t.Errorf("Header mismatch: %v != %v", h0, h)
		} else if bytes.Compare(h, empty) == 0 {
			t.Error("Header appears uninitialized")
		} else if uint64(len(h)) != headerSize {
			t.Errorf("Header size mismatch: %d != %d", len(h), headerSize)
		}
		var frame = make(DataFrame, source.frameBytes)
		var nextFrame uint64
		for f := 0; f < 6; f++ {
			var err error
			if _, err = source.Next(&nextFrame, &frame); err != nil {
				if err != io.EOF {
					t.Error(err)
				}
				break
			}
		}
		source.Done()
	}
	CloseAllSources()
}

func TestContentEqual(t *testing.T) {
	source := GetSource("/dev/urandom", &Config{
		SourceBuffer: 32,
		FrameBytes:   65536,
		CloseIdle:    true,
	})
	nConsumers := 10
	done := make(chan uint64, nConsumers)
	tab := crc64.MakeTable(crc64.ECMA)
	for i := 0; i < nConsumers; i++ {
		go func() {
			var hash uint64
			defer func() { done <- hash }()
			var frame = make(DataFrame, source.frameBytes)
			var nextFrame uint64
			for f := 0; f < 130; f++ {
				var err error
				if _, err = source.Next(&nextFrame, &frame); err != nil && err != io.EOF {
					t.Fatalf("source.Next(%d, frame): %s", nextFrame, err)
				}
				hash = crc64.Update(hash, tab, frame)
			}
		}()
	}
	h0 := <-done
	for i := 1; i < nConsumers; i++ {
		hi := <-done
		if h0 == 0 {
			h0 = hi
		} else if hi != h0 {
			t.Errorf("hash mismatch: h0=%x, h%d=%x", h0, i, hi)
		}
	}
	source.Done()
}

// 32s of CD audio, 1s per frame, 1000 consumers
func BenchmarkSource1KConsumers(b *testing.B) {
	benchSource(b, 1000, Config{
		SourceBuffer: 32,
		FrameBytes:   176800,
	})
}

// 320s of CD audio, 1s per frame, 10000 consumers
func BenchmarkSource10KConsumers(b *testing.B) {
	benchSource(b, 10000, Config{
		SourceBuffer: 320,
		FrameBytes:   176800,
	})
}

// 32s of 128kbps audio, 1s per frame, 10000 consumers
func BenchmarkSource128kbps10KConsumers(b *testing.B) {
	benchSource(b, 10000, Config{
		SourceBuffer: 32,
		FrameBytes:   128 * 1000 / 8,
	})
}

func BenchmarkSourceTinyBuffer(b *testing.B) {
	benchSource(b, 1, Config{
		SourceBuffer: 3,
		FrameBytes:   1,
	})
}

func BenchmarkSourceMediumBuffer(b *testing.B) {
	benchSource(b, 1, Config{
		SourceBuffer: 64,
		FrameBytes:   64,
	})
}

func BenchmarkSourceBigFrame(b *testing.B) {
	benchSource(b, 1, Config{
		SourceBuffer: 64,
		FrameBytes:   1048576,
	})
}

func BenchmarkSourceBigBuffer(b *testing.B) {
	benchSource(b, 1, Config{
		SourceBuffer: 1048576,
		FrameBytes:   64,
	})
}

func benchSource(b *testing.B, nConsumers int, c Config) {
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(runtime.NumCPU()))
	source := GetSource("/dev/zero", &c)
	wg := &sync.WaitGroup{}
	wg.Add(nConsumers)
	for c := 0; c < nConsumers; c++ {
		go func(c int) {
			defer wg.Done()
			consume(b, source, c)
		}(c)
	}
	wg.Wait()
	CloseAllSources()
}

func consume(b *testing.B, source *Source, label interface{}) {
	var frame DataFrame = make(DataFrame, source.frameBytes)
	var nextFrame uint64
	for i := uint64(0); i < 10*uint64(b.N); i++ {
		if _, err := source.Next(&nextFrame, &frame); err != nil {
			if err != io.EOF {
				b.Fatalf("source.Next(%d, frame): %s", nextFrame, err)
			}
			break
		}
	}
}
<commit_msg>Test Source's filtering feature.<commit_after>package main

import (
	"bytes"
	"fmt"
	"hash/crc64"
	"io"
	"math/rand"
	"os"
	"os/exec"
	"runtime"
	"sync"
	"testing"
	"time"
)

func TestSigpipe(t *testing.T) {
	r, w, err := os.Pipe()
	if err != nil {
		t.Fatal(err)
	}
	cmd := exec.Command("head", "-c64", "/dev/urandom")
	cmd.Stdout = w
	if err = cmd.Start(); err != nil {
		t.Fatal(err)
	}
	w.Close()
	go cmd.Wait()
	source := GetSource(fmt.Sprintf("/dev/fd/%d", r.Fd()), &Config{
		SourceBuffer: 5,
		FrameBytes:   16,
		CloseIdle:    false,
		Reopen:       true,
	})
	var frame = make(DataFrame, source.frameBytes)
	var nextFrame uint64
	for f := 0; f < 4; f++ {
		if _, err = source.Next(&nextFrame, &frame); err != nil {
			t.Error(err)
			break
		}
	}
	done := make(chan error, 1)
	go func() {
		_, err = source.Next(&nextFrame, &frame)
		done <- err
	}()
	select {
	case <-time.After(time.Second):
		t.Error("Should have given up within 1s of SIGCHLD")
	case err = <-done:
		if err == nil {
			t.Errorf("Should have error, got frame %v", frame)
		}
	}
	r.Close()
	source.Done()
	CloseAllSources()
}

func TestEmptySource(t *testing.T) {
	r, w, err := os.Pipe()
	if err != nil {
		t.Fatal(err)
	}
	closeNow := make(chan bool)
	go func() {
		<-closeNow
		w.Close()
	}()
	source := GetSource(fmt.Sprintf("/dev/fd/%d", r.Fd()), &Config{
		SourceBuffer: 5,
		FrameBytes:   16,
		CloseIdle:    false,
		Reopen:       false,
	})
	var frame = make(DataFrame, source.frameBytes)
	var nextFrame uint64
	done := make(chan error, 1)
	go func() {
		_, err := source.Next(&nextFrame, &frame)
		done <- err
	}()
	time.Sleep(time.Millisecond)
	select {
	case err := <-done:
		t.Errorf("Should still be waiting for input, got error %v", err)
	default:
	}
	closeNow <- true
	select {
	case <-time.After(100 * time.Millisecond):
		t.Error("Should have given up within 100ms")
	case err := <-done:
		if err == nil {
			t.Errorf("Should have error, got frame %v", frame)
		}
	}
	source.Done()
	r.Close()
	CloseAllSources()
}

func failUnless(t *testing.T, ms time.Duration, ok <-chan bool) {
	select {
	case <-time.After(ms * time.Millisecond):
		t.Fatalf("Timed out in %d ms", ms)
	case <-ok:
	}
}

func DataFaker(t *testing.T) (string, chan<- interface{}, *[]byte) {
	r, w, err := os.Pipe()
	if err != nil {
		t.Fatal(err)
	}
	var sentData []byte
	wantMore := make(chan interface{})
	go func() {
		for cmd := range wantMore {
			var moreData []byte
			var n int
			if n, ok := cmd.(int); ok {
				moreData = make([]byte, n)
			} else {
				moreData = cmd.([]byte)
			}
			for i := 0; i < n; i++ {
				moreData[i] = byte(rand.Int() & 0xff)
			}
			sentData = append(sentData, moreData...)
			if did, err := w.Write(moreData); did < len(moreData) || err != nil {
				t.Error("Short write: %d < %d, %s", did, len(moreData), err)
			}
			w.Sync()
		}
		r.Close()
		w.Close()
	}()
	return fmt.Sprintf("/dev/fd/%d", r.Fd()), wantMore, &sentData
}

func TestSourceFilter(t *testing.T) {
	type expect struct {
		frame     []byte
		frameSize int
		err       error
	}
	fMock := make(chan *expect, 100)
	var pending *expect
	Filters["MOCK"] = func(frame []byte, context *interface{}) (frameSize int, err error) {
		if pending == nil {
			pending = <-fMock
		}
		if len(frame) < len(pending.frame) {
			return len(pending.frame), ShortFrame
		}
		if 0 != bytes.Compare(pending.frame, frame[0:len(pending.frame)]) {
			t.Fatalf("Expected filter(%v), got %v", pending.frame, frame)
		}
		defer func() { pending = nil }()
		return pending.frameSize, pending.err
	}
	defer func() { delete(Filters, "MOCK") }()
	fakeFile, sendFake, _ := DataFaker(t)
	source := GetSource(fakeFile, &Config{
		SourceBuffer: 5,
		FrameBytes:   4,
		CloseIdle:    true,
		Reopen:       false,
		FrameFilter:  "MOCK",
	})
	var nextFrame uint64
	var frame = make(DataFrame, 4)
	expectNext := func(expect []byte, expectErr error) {
		_, err := source.Next(&nextFrame, &frame)
		if err != expectErr {
			t.Errorf("Frame %d expected err %s got %s", nextFrame, expectErr, err)
		} else if err == nil && 0 != bytes.Compare(expect, frame) {
			t.Errorf("Frame %d expected bytes %v got %v", nextFrame, expect, frame)
		}
	}
	// Our mock filter will pass any frame without a 00.
	sendFake <- []byte{11, 22, 33, 44, 11, 22, 33, 00, 11, 00, 33, 44, 11, 22, 33, 44}
	fMock <- &expect{[]byte{11, 22, 33, 44}, 4, nil}
	expectNext([]byte{11, 22, 33, 44}, nil)
	fMock <- &expect{[]byte{11, 22, 33, 00}, 3, nil}
	expectNext([]byte{11, 22, 33}, nil)
	fMock <- &expect{[]byte{00, 11, 00, 33}, 0, InvalidFrame}
	fMock <- &expect{[]byte{11, 00, 33, 44}, 1, nil}
	expectNext([]byte{11}, nil)
	fMock <- &expect{[]byte{00, 33, 44, 11}, 0, InvalidFrame}
	fMock <- &expect{[]byte{33, 44, 11, 22}, 4, nil}
	expectNext([]byte{33, 44, 11, 22}, nil)
	source.Done()
	// TODO: source should give filter a chance to accept the data at EOF, even though it doesn't fill max frame size.
	// fMock <- &expect{[]byte{33,44}, 1, nil}
	// fMock <- &expect{[]byte{44}, 1, nil}
	// expectNext([]byte{33}, nil)
	// expectNext([]byte{44}, nil)
	expectNext([]byte{}, io.EOF)
}

func TestSentEqualsReceived(t *testing.T) {
	fakeData, wantMore, sentData := DataFaker(t)
	source := GetSource(fakeData, &Config{
		SourceBuffer: 5,
		FrameBytes:   16,
		CloseIdle:    true,
		Reopen:       false,
	})
	var rcvdData = []byte{}
	var frame = make(DataFrame, source.frameBytes)
	var nextFrame uint64
	ok := make(chan bool)
	wantMore <- 19 // Send one full frame (to make sure we start receiving at frame 0) and a bit extra
	for f := 100; f > 0; f-- {
		go func() {
			if _, err := source.Next(&nextFrame, &frame); err != nil {
				t.Fatal(err)
			}
			if len(frame) != 16 {
				t.Errorf("Wrong size frame, want 16 got %d", len(frame))
			}
			ok <- true
		}()
		failUnless(t, 100, ok)
		rcvdData = append(rcvdData, frame...)
		if f == 4 {
			wantMore <- 16 * 3 // Last three frames (we sent one frame before the loop)
		} else if f%4 == 0 {
			wantMore <- 16 * 4 // Next four frames
		}
	}
	close(wantMore)
	go func() {
		if _, err := source.Next(&nextFrame, &frame); err != io.EOF {
			t.Error("Should have got EOF")
		}
		ok <- true
	}()
	failUnless(t, 100, ok)
	wantData := (*sentData)[0:1600]
	if 0 != bytes.Compare(wantData, rcvdData) {
		t.Errorf("want %d != rcvd %d", len(wantData), len(rcvdData))
	}
	source.Done()
	CloseAllSources()
}

func TestHeader(t *testing.T) {
	headerSize := uint64(64)
	nClients := 5
	sources := make(chan *Source, nClients)
	for i := 0; i < nClients; i++ {
		sources <- GetSource("/dev/urandom", &Config{
			SourceBuffer: 5,
			FrameBytes:   65536,
			HeaderBytes:  headerSize,
			CloseIdle:    true,
		})
	}
	empty := make([]byte, headerSize)
	var h0 []byte
	for i := 0; i < nClients; i++ {
		h := make([]byte, headerSize)
		source := <-sources
		err := source.GetHeader(h)
		if err != nil {
			t.Error(err)
		} else if h0 == nil {
			h0 = h
		} else if bytes.Compare(h, h0) != 0 {
			t.Errorf("Header mismatch: %v != %v", h0, h)
		} else if bytes.Compare(h, empty) == 0 {
			t.Error("Header appears uninitialized")
		} else if uint64(len(h)) != headerSize {
			t.Errorf("Header size mismatch: %d != %d", len(h), headerSize)
		}
		var frame = make(DataFrame, source.frameBytes)
		var nextFrame uint64
		for f := 0; f < 6; f++ {
			var err error
			if _, err = source.Next(&nextFrame, &frame); err != nil {
				if err != io.EOF {
					t.Error(err)
				}
				break
			}
		}
		source.Done()
	}
	CloseAllSources()
}

func TestContentEqual(t *testing.T) {
	source := GetSource("/dev/urandom", &Config{
		SourceBuffer: 32,
		FrameBytes:   65536,
		CloseIdle:    true,
	})
	nConsumers := 10
	done := make(chan uint64, nConsumers)
	tab := crc64.MakeTable(crc64.ECMA)
	for i := 0; i < nConsumers; i++ {
		go func() {
			var hash uint64
			defer func() { done <- hash }()
			var frame = make(DataFrame, source.frameBytes)
			var nextFrame uint64
			for f := 0; f < 130; f++ {
				var err error
				if _, err = source.Next(&nextFrame, &frame); err != nil && err != io.EOF {
					t.Fatalf("source.Next(%d, frame): %s", nextFrame, err)
				}
				hash = crc64.Update(hash, tab, frame)
			}
		}()
	}
	h0 := <-done
	for i := 1; i < nConsumers; i++ {
		hi := <-done
		if h0 == 0 {
			h0 = hi
		} else if hi != h0 {
			t.Errorf("hash mismatch: h0=%x, h%d=%x", h0, i, hi)
		}
	}
	source.Done()
}

// 32s of CD audio, 1s per frame, 1000 consumers
func BenchmarkSource1KConsumers(b *testing.B) {
	benchSource(b, 1000, Config{
		SourceBuffer: 32,
		FrameBytes:   176800,
	})
}

// 320s of CD audio, 1s per frame, 10000 consumers
func BenchmarkSource10KConsumers(b *testing.B) {
	benchSource(b, 10000, Config{
		SourceBuffer: 320,
		FrameBytes:   176800,
	})
}

// 32s of 128kbps audio, 1s per frame, 10000 consumers
func BenchmarkSource128kbps10KConsumers(b *testing.B) {
	benchSource(b, 10000, Config{
		SourceBuffer: 32,
		FrameBytes:   128 * 1000 / 8,
	})
}

func BenchmarkSourceTinyBuffer(b *testing.B) {
	benchSource(b, 1, Config{
		SourceBuffer: 3,
		FrameBytes:   1,
	})
}

func BenchmarkSourceMediumBuffer(b *testing.B) {
	benchSource(b, 1, Config{
		SourceBuffer: 64,
		FrameBytes:   64,
	})
}

func BenchmarkSourceBigFrame(b *testing.B) {
	benchSource(b, 1, Config{
		SourceBuffer: 64,
		FrameBytes:   1048576,
	})
}

func BenchmarkSourceBigBuffer(b *testing.B) {
	benchSource(b, 1, Config{
		SourceBuffer: 1048576,
		FrameBytes:   64,
	})
}

func benchSource(b *testing.B, nConsumers int, c Config) {
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(runtime.NumCPU()))
	source := GetSource("/dev/zero", &c)
	wg := &sync.WaitGroup{}
	wg.Add(nConsumers)
	for c := 0; c < nConsumers; c++ {
		go func(c int) {
			defer wg.Done()
			consume(b, source, c)
		}(c)
	}
	wg.Wait()
	CloseAllSources()
}

func consume(b *testing.B, source *Source, label interface{}) {
	var frame DataFrame = make(DataFrame, source.frameBytes)
	var nextFrame uint64
	for i := uint64(0); i < 10*uint64(b.N); i++ {
		if _, err := source.Next(&nextFrame, &frame); err != nil {
			if err != io.EOF {
				b.Fatalf("source.Next(%d, frame): %s", nextFrame, err)
			}
			break
		}
	}
}
<|endoftext|>


<commit_before>/*
Tensile web stress test tool

Mike Hughes 2014
intermernet AT gmail DOT com

LICENSE BSD 3 Clause
*/
package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"runtime"
	"sync"
	"time"
)

const version = "0.1"

var (
	url string

	reqs int
	max  int

	wg sync.WaitGroup
)

func init() {
	flag.StringVar(&url, "url", "http://localhost/", "Target URL")
	flag.IntVar(&reqs, "reqs", 1000, "Total requests")
	flag.IntVar(&max, "concurrent", 100, "Maximum concurrent requests")
}

type Response struct {
	*http.Response
	err error
}

// Dispatcher
func dispatcher(reqChan chan *http.Request) {
	defer close(reqChan)
	for i := 0; i < reqs; i++ {
		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			log.Println(err)
		}
		reqChan <- req
	}
}

// Worker Pool
func workerPool(reqChan chan *http.Request, respChan chan Response) {
	defer close(respChan)
	t := &http.Transport{}
	defer t.CloseIdleConnections()
	for i := 0; i < max; i++ {
		wg.Add(1)
		go worker(t, reqChan, respChan)
	}
	wg.Wait()
}

// Worker
func worker(t *http.Transport, reqChan chan *http.Request, respChan chan Response) {
	for req := range reqChan {
		resp, err := t.RoundTrip(req)
		r := Response{resp, err}
		respChan <- r
	}
	wg.Done()
}

// Consumer
func consumer(respChan chan Response) (int64, int64) {
	var (
		conns int64
		size  int64
	)
	for r := range respChan {
		if r.err != nil {
			log.Println(r.err)
		} else {
			size += r.ContentLength
			if err := r.Body.Close(); err != nil {
				log.Println(r.err)
			}
		}
		conns++
	}
	return conns, size
}

func main() {
	flag.Parse()
	fmt.Printf("\n\tTensile web stress test tool v%s\n\n", version)
	if reqs < max {
		fmt.Println("NOTICE: Concurrent requests is greater than number of requests.")
		fmt.Println("\tChanging concurrent requests to number of requests\n")
		max = reqs
	}
	runtime.GOMAXPROCS(runtime.NumCPU())
	reqChan := make(chan *http.Request)
	respChan := make(chan Response)
	fmt.Printf("Sending %d requests to %s with %d concurrent workers.\n\n", reqs, url, max)
	start := time.Now()
	go dispatcher(reqChan)
	go workerPool(reqChan, respChan)
	fmt.Println("Waiting for replies...\n")
	conns, size := consumer(respChan)
	took := time.Since(start)
	ns := took.Nanoseconds()
	var av int64
	if conns != 0 {
		av = ns / conns
	}
	average, err := time.ParseDuration(fmt.Sprintf("%d", av) + "ns")
	if err != nil {
		log.Println(err)
	}
	fmt.Printf("Connections:\t%d\nConcurrent:\t%d\nTotal size:\t%d bytes\nTotal time:\t%s\nAverage time:\t%s\n", conns, max, size, took, average)
}
<commit_msg>Defer WaitGroup Done() in worker<commit_after>/*
Tensile web stress test tool

Mike Hughes 2014
intermernet AT gmail DOT com

LICENSE BSD 3 Clause
*/
package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"runtime"
	"sync"
	"time"
)

const version = "0.1"

var (
	url string

	reqs int
	max  int

	wg sync.WaitGroup
)

func init() {
	flag.StringVar(&url, "url", "http://localhost/", "Target URL")
	flag.IntVar(&reqs, "reqs", 1000, "Total requests")
	flag.IntVar(&max, "concurrent", 100, "Maximum concurrent requests")
}

type Response struct {
	*http.Response
	err error
}

// Dispatcher
func dispatcher(reqChan chan *http.Request) {
	defer close(reqChan)
	for i := 0; i < reqs; i++ {
		req, err := http.NewRequest("GET", url, nil)
		if err != nil {
			log.Println(err)
		}
		reqChan <- req
	}
}

// Worker Pool
func workerPool(reqChan chan *http.Request, respChan chan Response) {
	defer close(respChan)
	t := &http.Transport{}
	defer t.CloseIdleConnections()
	for i := 0; i < max; i++ {
		wg.Add(1)
		go worker(t, reqChan, respChan)
	}
	wg.Wait()
}

// Worker
func worker(t *http.Transport, reqChan chan *http.Request, respChan chan Response) {
	defer wg.Done()
	for req := range reqChan {
		resp, err := t.RoundTrip(req)
		r := Response{resp, err}
		respChan <- r
	}
}

// Consumer
func consumer(respChan chan Response) (int64, int64) {
	var (
		conns int64
		size  int64
	)
	for r := range respChan {
		if r.err != nil {
			log.Println(r.err)
		} else {
			size += r.ContentLength
			if err := r.Body.Close(); err != nil {
				log.Println(r.err)
			}
		}
		conns++
	}
	return conns, size
}

func main() {
	flag.Parse()
	fmt.Printf("\n\tTensile web stress test tool v%s\n\n", version)
	if reqs < max {
		fmt.Println("NOTICE: Concurrent requests is greater than number of requests.")
		fmt.Println("\tChanging concurrent requests to number of requests\n")
		max = reqs
	}
	runtime.GOMAXPROCS(runtime.NumCPU())
	reqChan := make(chan *http.Request)
	respChan := make(chan Response)
	fmt.Printf("Sending %d requests to %s with %d concurrent workers.\n\n", reqs, url, max)
	start := time.Now()
	go dispatcher(reqChan)
	go workerPool(reqChan, respChan)
	fmt.Println("Waiting for replies...\n")
	conns, size := consumer(respChan)
	took := time.Since(start)
	ns := took.Nanoseconds()
	var av int64
	if conns != 0 {
		av = ns / conns
	}
	average, err := time.ParseDuration(fmt.Sprintf("%d", av) + "ns")
	if err != nil {
		log.Println(err)
	}
	fmt.Printf("Connections:\t%d\nConcurrent:\t%d\nTotal size:\t%d bytes\nTotal time:\t%s\nAverage time:\t%s\n", conns, max, size, took, average)
}
<|endoftext|>


<commit_before>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package stringforwarder_test

import (
	"time"

	coretesting "github.com/juju/juju/testing"
	"github.com/juju/juju/utils/stringforwarder"

	gc "gopkg.in/check.v1"
)

type stringForwarderSuite struct{}

var _ = gc.Suite(&stringForwarderSuite{})

func (*stringForwarderSuite) TestReceives(c *gc.C) {
	messages := make([]string, 0)
	received := make(chan struct{}, 10)
	forwarder := stringforwarder.NewStringForwarder(func(msg string) {
		messages = append(messages, msg)
		received <- struct{}{}
	})
	forwarder.Receive("one")
	select {
	case <-received:
	case <-time.After(coretesting.LongWait):
		c.Errorf("timeout waiting for a message to be received")
	}
	c.Check(forwarder.Stop(), gc.Equals, 0)
	c.Check(messages, gc.DeepEquals, []string{"one"})
}

func noopCallback(string) {
}

func (*stringForwarderSuite) TestStopIsReentrant(c *gc.C) {
	forwarder := stringforwarder.NewStringForwarder(noopCallback)
	forwarder.Stop()
	forwarder.Stop()
}

func (*stringForwarderSuite) TestMessagesDroppedAfterStop(c *gc.C) {
	messages := make([]string, 0)
	forwarder := stringforwarder.NewStringForwarder(func(msg string) {
		messages = append(messages, msg)
	})
	forwarder.Stop()
	forwarder.Receive("one")
	forwarder.Receive("two")
	forwarder.Stop()
	c.Check(messages, gc.DeepEquals, []string{})
}

func (*stringForwarderSuite) TestAllDroppedWithNoCallback(c *gc.C) {
	forwarder := stringforwarder.NewStringForwarder(nil)
	forwarder.Receive("one")
	forwarder.Receive("two")
	forwarder.Receive("three")
	c.Check(forwarder.Stop(), gc.Equals, 3)
}

func (*stringForwarderSuite) TestMessagesDroppedWhenBusy(c *gc.C) {
	messages := make([]string, 0)
	next := make(chan struct{})
	blockingCallback := func(msg string) {
		select {
		case <-next:
			messages = append(messages, msg)
		case <-time.After(coretesting.LongWait):
			c.Error("timeout waiting for next")
			return
		}
	}
	forwarder := stringforwarder.NewStringForwarder(blockingCallback)
	forwarder.Receive("first")
	forwarder.Receive("second")
	forwarder.Receive("third")
	// At this point we should have started processing "first", but the
	// other two messages are dropped.
	select {
	case next <- struct{}{}:
	case <-time.After(coretesting.LongWait):
		c.Error("failed to send next signal")
	}
	// now we should be ready to get another message
	forwarder.Receive("fourth")
	forwarder.Receive("fifth")
	// finish fourth
	select {
	case next <- struct{}{}:
	case <-time.After(coretesting.LongWait):
		c.Error("failed to send next signal")
	}
	dropCount := forwarder.Stop()
	c.Check(messages, gc.DeepEquals, []string{"first", "fourth"})
	c.Check(dropCount, gc.Equals, 3)
}
<commit_msg>fix import order<commit_after>// Copyright 2016 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package stringforwarder_test

import (
	"time"

	gc "gopkg.in/check.v1"

	coretesting "github.com/juju/juju/testing"
	"github.com/juju/juju/utils/stringforwarder"
)

type stringForwarderSuite struct{}

var _ = gc.Suite(&stringForwarderSuite{})

func (*stringForwarderSuite) TestReceives(c *gc.C) {
	messages := make([]string, 0)
	received := make(chan struct{}, 10)
	forwarder := stringforwarder.NewStringForwarder(func(msg string) {
		messages = append(messages, msg)
		received <- struct{}{}
	})
	forwarder.Receive("one")
	select {
	case <-received:
	case <-time.After(coretesting.LongWait):
		c.Errorf("timeout waiting for a message to be received")
	}
	c.Check(forwarder.Stop(), gc.Equals, 0)
	c.Check(messages, gc.DeepEquals, []string{"one"})
}

func noopCallback(string) {
}

func (*stringForwarderSuite) TestStopIsReentrant(c *gc.C) {
	forwarder := stringforwarder.NewStringForwarder(noopCallback)
	forwarder.Stop()
	forwarder.Stop()
}

func (*stringForwarderSuite) TestMessagesDroppedAfterStop(c *gc.C) {
	messages := make([]string, 0)
	forwarder := stringforwarder.NewStringForwarder(func(msg string) {
		messages = append(messages, msg)
	})
	forwarder.Stop()
	forwarder.Receive("one")
	forwarder.Receive("two")
	forwarder.Stop()
	c.Check(messages, gc.DeepEquals, []string{})
}

func (*stringForwarderSuite) TestAllDroppedWithNoCallback(c *gc.C) {
	forwarder := stringforwarder.NewStringForwarder(nil)
	forwarder.Receive("one")
	forwarder.Receive("two")
	forwarder.Receive("three")
	c.Check(forwarder.Stop(), gc.Equals, 3)
}

func (*stringForwarderSuite) TestMessagesDroppedWhenBusy(c *gc.C) {
	messages := make([]string, 0)
	next := make(chan struct{})
	blockingCallback := func(msg string) {
		select {
		case <-next:
			messages = append(messages, msg)
		case <-time.After(coretesting.LongWait):
			c.Error("timeout waiting for next")
			return
		}
	}
	forwarder := stringforwarder.NewStringForwarder(blockingCallback)
	forwarder.Receive("first")
	forwarder.Receive("second")
	forwarder.Receive("third")
	// At this point we should have started processing "first", but the
	// other two messages are dropped.
	select {
	case next <- struct{}{}:
	case <-time.After(coretesting.LongWait):
		c.Error("failed to send next signal")
	}
	// now we should be ready to get another message
	forwarder.Receive("fourth")
	forwarder.Receive("fifth")
	// finish fourth
	select {
	case next <- struct{}{}:
	case <-time.After(coretesting.LongWait):
		c.Error("failed to send next signal")
	}
	dropCount := forwarder.Stop()
	c.Check(messages, gc.DeepEquals, []string{"first", "fourth"})
	c.Check(dropCount, gc.Equals, 3)
}
<|endoftext|>


<commit_before>package sphinx

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/Clever/sphinx/common"
	"github.com/Clever/sphinx/matchers"
	"strings"
	"testing"
)

func checkStatusForRequests(ratelimiter RateLimiter,
	request common.Request, num int, expectedStatuses []Status) error {

	var statuses []Status
	var err error
	for i := 0; i < num; i++ {
		statuses, err = ratelimiter.Add(request)
		if err != nil {
			return err
		}
	}

	if len(statuses) != len(expectedStatuses) {
		return fmt.Errorf("expected to match %d buckets. Got: %d",
			len(expectedStatuses), len(statuses))
	}
	for i, status := range expectedStatuses {
		if status.Remaining != statuses[i].Remaining && status.Name != statuses[i].Name {
			return errors.New("expected 5 requests for the 'bearer/events' limits")
		}
	}

	return nil
}

// ratelimiter is initialized properly based on config
func TestNewRateLimiter(t *testing.T) {

	config, err := NewConfiguration("./example.yaml")
	if err != nil {
		t.Error("could not load example configuration")
	}

	ratelimiter, err := NewRateLimiter(config)
	if err != nil {
		t.Errorf("Error while instantiating ratelimiter: %s", err.Error())
	}
	if len(ratelimiter.Configuration().Limits) !=
		len(ratelimiter.Limits()) {
		t.Error("expected number of limits in configuration to match instantiated limits")
	}
}

// test that matcher errors are bubbled up
func TestBadConfiguration(t *testing.T) {

	var configBuf = bytes.NewBufferString(`
proxy:
  handler: http
  host: http://proxy.example.com
  listen: :8080
storage:
  type: memory
limits:
  test:
    interval: 15  # in seconds
    max: 200
`)

	// header matchers are verified
	configBuf.WriteString(`
    matches:
      headers:
        match_any:
          - "Authorization": "Bearer.*"
          - name: "X-Forwarded-For"
`)
	configuration, err := loadAndValidateConfig(configBuf.Bytes())
	if err != nil {
		t.Error("configuration failed with error", err)
	}

	_, err = NewRateLimiter(configuration)
	if err == nil {
		t.Error("Expected header matcher error, got none")
	} else if !strings.Contains(err.Error(), "InvalidMatcherConfig: headers") {
		t.Errorf("Expected a InvalidMatcherConfig error, got different error: %s", err.Error())
	}

}

// adds different kinds of requests and checks limit Status
// focusses on single bucket adds
func TestSimpleAdd(t *testing.T) {
	config, err := NewConfiguration("./example.yaml")
	if err != nil {
		t.Error("could not load example configuration")
	}
	ratelimiter, err := NewRateLimiter(config)

	request := common.Request{
		"path": "/v1.1/events/students/123",
		"headers": common.ConstructMockRequestWithHeaders(map[string][]string{
			"Authorization":   []string{"Bearer 12345"},
			"X-Forwarded-For": []string{"IP1", "IP2"},
		}).Header,
		"remoteaddr": "127.0.0.1",
	}
	if err = checkStatusForRequests(
		ratelimiter, request, 5, []Status{
			Status{Remaining: 195, Name: "bearer/events"}}); err != nil {
		t.Error(err)
	}

	request = common.Request{
		"path": "/v1.1/students/123",
		"headers": common.ConstructMockRequestWithHeaders(map[string][]string{
			"Authorization": []string{"Basic 12345"},
		}).Header,
	}

	if err = checkStatusForRequests(
		ratelimiter, request, 1, []Status{
			Status{Remaining: 195, Name: "basic/main"}}); err != nil {
		t.Error(err)
	}
}

type NeverMatch struct{}

func (m NeverMatch) Match(req common.Request) bool {
	return false
}

func createLimit(numMatchers int) *Limit {
	neverMatchers := []matchers.Matcher{}
	for i := 0; i < numMatchers; i++ {
		neverMatchers = append(neverMatchers, NeverMatch{})
	}
	limit := &Limit{
		matcher: requestMatcher{
			Matches:  neverMatchers,
			Excludes: neverMatchers,
		},
	}
	return limit
}

var benchMatch = func(b *testing.B, numMatchers int) {
	limit := createLimit(numMatchers)
	request := common.Request{}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		limit.match(request)
	}
}

func BenchmarkMatch1(b *testing.B) {
	benchMatch(b, 1)
}

func BenchmarkMatch100(b *testing.B) {
	benchMatch(b, 100)
}
<commit_msg>sphinx: benchmark RateLimiter.Add<commit_after>package sphinx

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/Clever/sphinx/common"
	"github.com/Clever/sphinx/matchers"
	"strings"
	"testing"
)

func checkStatusForRequests(ratelimiter RateLimiter,
	request common.Request, num int, expectedStatuses []Status) error {

	var statuses []Status
	var err error
	for i := 0; i < num; i++ {
		statuses, err = ratelimiter.Add(request)
		if err != nil {
			return err
		}
	}

	if len(statuses) != len(expectedStatuses) {
		return fmt.Errorf("expected to match %d buckets. Got: %d",
			len(expectedStatuses), len(statuses))
	}
	for i, status := range expectedStatuses {
		if status.Remaining != statuses[i].Remaining && status.Name != statuses[i].Name {
			return errors.New("expected 5 requests for the 'bearer/events' limits")
		}
	}

	return nil
}

// ratelimiter is initialized properly based on config
func TestNewRateLimiter(t *testing.T) {

	config, err := NewConfiguration("./example.yaml")
	if err != nil {
		t.Error("could not load example configuration")
	}

	ratelimiter, err := NewRateLimiter(config)
	if err != nil {
		t.Errorf("Error while instantiating ratelimiter: %s", err.Error())
	}
	if len(ratelimiter.Configuration().Limits) !=
		len(ratelimiter.Limits()) {
		t.Error("expected number of limits in configuration to match instantiated limits")
	}
}

// test that matcher errors are bubbled up
func TestBadConfiguration(t *testing.T) {

	var configBuf = bytes.NewBufferString(`
proxy:
  handler: http
  host: http://proxy.example.com
  listen: :8080
storage:
  type: memory
limits:
  test:
    interval: 15  # in seconds
    max: 200
`)

	// header matchers are verified
	configBuf.WriteString(`
    matches:
      headers:
        match_any:
          - "Authorization": "Bearer.*"
          - name: "X-Forwarded-For"
`)
	configuration, err := loadAndValidateConfig(configBuf.Bytes())
	if err != nil {
		t.Error("configuration failed with error", err)
	}

	_, err = NewRateLimiter(configuration)
	if err == nil {
		t.Error("Expected header matcher error, got none")
	} else if !strings.Contains(err.Error(), "InvalidMatcherConfig: headers") {
		t.Errorf("Expected a InvalidMatcherConfig error, got different error: %s", err.Error())
	}

}

// adds different kinds of requests and checks limit Status
// focusses on single bucket adds
func TestSimpleAdd(t *testing.T) {
	config, err := NewConfiguration("./example.yaml")
	if err != nil {
		t.Error("could not load example configuration")
	}
	ratelimiter, err := NewRateLimiter(config)

	request := common.Request{
		"path": "/v1.1/events/students/123",
		"headers": common.ConstructMockRequestWithHeaders(map[string][]string{
			"Authorization":   []string{"Bearer 12345"},
			"X-Forwarded-For": []string{"IP1", "IP2"},
		}).Header,
		"remoteaddr": "127.0.0.1",
	}
	if err = checkStatusForRequests(
		ratelimiter, request, 5, []Status{
			Status{Remaining: 195, Name: "bearer/events"}}); err != nil {
		t.Error(err)
	}

	request = common.Request{
		"path": "/v1.1/students/123",
		"headers": common.ConstructMockRequestWithHeaders(map[string][]string{
			"Authorization": []string{"Basic 12345"},
		}).Header,
	}

	if err = checkStatusForRequests(
		ratelimiter, request, 1, []Status{
			Status{Remaining: 195, Name: "basic/main"}}); err != nil {
		t.Error(err)
	}
}

type NeverMatch struct{}

func (m NeverMatch) Match(req common.Request) bool {
	return false
}

func createLimit(numMatchers int) *Limit {
	neverMatchers := []matchers.Matcher{}
	for i := 0; i < numMatchers; i++ {
		neverMatchers = append(neverMatchers, NeverMatch{})
	}
	limit := &Limit{
		matcher: requestMatcher{
			Matches:  neverMatchers,
			Excludes: neverMatchers,
		},
	}
	return limit
}

var benchMatch = func(b *testing.B, numMatchers int) {
	limit := createLimit(numMatchers)
	request := common.Request{}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		limit.match(request)
	}
}

func BenchmarkMatch1(b *testing.B) {
	benchMatch(b, 1)
}

func BenchmarkMatch100(b *testing.B) {
	benchMatch(b, 100)
}

func createRateLimiter(numLimits int) RateLimiter {
	limit := createLimit(1)
	rateLimiter := &sphinxRateLimiter{}
	limits := []*Limit{}
	for i := 0; i < numLimits; i++ {
		limits = append(limits, limit)
	}
	rateLimiter.limits = limits
	return rateLimiter
}

var benchAdd = func(b *testing.B, numLimits int) {
	rateLimiter := createRateLimiter(numLimits)
	request := common.Request{}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		rateLimiter.Add(request)
	}
}

func BenchmarkAdd1(b *testing.B) {
	benchAdd(b, 1)
}

func BenchmarkAdd100(b *testing.B) {
	benchAdd(b, 100)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
)

const (
	HELP = `md2xml [-h] file.md
Transform a given Markdown file into XML.
-h        To print this help page.
-x        Print intermediate XHTML output.
-a        Output article (instead of blog entry).
-s        Print stylesheet used for transformation.
file.md   The markdown file to convert.
Note: this program calls xsltproc that must have been installed.`
	STYLESHEET_ARTICLE = `<?xml version="1.0" encoding="utf-8"?>

<xsl:stylesheet xmlns:xsl="http://www.w3.org/1999/XSL/Transform"
                version="1.0">

  <xsl:output method="xml" encoding="UTF-8"/>
  <xsl:param name="id">ID</xsl:param>
  <xsl:param name="date">DATE</xsl:param>
  <xsl:param name="title">TITLE</xsl:param>
  <xsl:param name="author">AUTHOR</xsl:param>
  <xsl:param name="email">EMAIL</xsl:param>

  <xsl:template match="/xhtml">
    <xsl:text disable-output-escaping="yes">
    <!DOCTYPE article PUBLIC "-//CAFEBABE//DTD blog 1.0//EN"
                             "../dtd/article.dtd">
    </xsl:text>

      <xsl:attribute name="id"><xsl:value-of select="$id"/></xsl:attribute>
      <xsl:attribute name="date"><xsl:value-of select="$date"/></xsl:attribute>
      <xsl:attribute name="author"><xsl:value-of select="$author"/></xsl:attribute>
      <xsl:attribute name="email"><xsl:value-of select="$email"/></xsl:attribute>
      <title><xsl:value-of select="$title"/></title>
      <text>
       <xsl:apply-templates/>
      </text>

  </xsl:template>

  <xsl:template match="h1">
    <sect level="1"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="h2">
    <sect level="2"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="h3">
    <sect level="3"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="h4">
    <sect level="4"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="h5">
    <sect level="5"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="h6">
    <sect level="6"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="p[@class='caption']">
  </xsl:template>

  <xsl:template match="p[count(text())=0 and count(code)=1]">
    <source><xsl:apply-templates select="code"/></source>
  </xsl:template>

  <xsl:template match="p[count(text())=1 and count(img)=1]">
    <xsl:apply-templates select="img"/>
  </xsl:template>

  <xsl:template match="img">

      <xsl:if test="@title">
        <title><xsl:value-of select="@title"/></title>
      </xsl:if>

  </xsl:template>

  <xsl:template match="p">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="ul">
    <list><xsl:apply-templates/></list>
  </xsl:template>

  <xsl:template match="ol">
    <enum><xsl:apply-templates/></enum>
  </xsl:template>

  <xsl:template match="li">
    <item><xsl:apply-templates/></item>
  </xsl:template>

  <xsl:template match="table">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="th">
    <xsl:apply-templates/>
  </xsl:template>

  <xsl:template match="tr">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="td">
    <co><xsl:apply-templates/></co>
  </xsl:template>

  <xsl:template match="pre">
    <source><xsl:apply-templates/></source>
  </xsl:template>

  <xsl:template match="em">
    <term><xsl:apply-templates/></term>
  </xsl:template>

  <xsl:template match="strong">
    <imp><xsl:apply-templates/></imp>
  </xsl:template>

  <xsl:template match="a">
    <link url="{@href}"><xsl:apply-templates/></link>
  </xsl:template>

</xsl:stylesheet>`
	STYLESHEET_BLOG = `<?xml version="1.0" encoding="utf-8"?>

<xsl:stylesheet xmlns:xsl="http://www.w3.org/1999/XSL/Transform"
                version="1.0">

  <xsl:output method="xml" encoding="UTF-8"/>
  <xsl:param name="id">ID</xsl:param>
  <xsl:param name="date">DATE</xsl:param>
  <xsl:param name="title">TITLE</xsl:param>

  <xsl:template match="/xhtml">
    <xsl:text disable-output-escaping="yes">
    <!DOCTYPE blog PUBLIC "-//CAFEBABE//DTD blog 1.0//EN"
                          "../dtd/blog.dtd">
    </xsl:text>
    <blog>
      <xsl:attribute name="id"><xsl:value-of select="$id"/></xsl:attribute>
      <xsl:attribute name="date"><xsl:value-of select="$date"/></xsl:attribute>
      <title><xsl:value-of select="$title"/></title>
      <xsl:apply-templates/>
    </blog>
  </xsl:template>

  <xsl:template match="h1">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="h2">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="h3">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="h4">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="h5">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="h6">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="p[@class='caption']">
  </xsl:template>

  <xsl:template match="p[count(text())=0 and count(code)=1]">
    <source><xsl:apply-templates select="code"/></source>
  </xsl:template>

  <xsl:template match="p[count(text())=1 and count(img)=1]">
    <xsl:apply-templates select="img"/>
  </xsl:template>

  <xsl:template match="img">

      <xsl:if test="@title">
        <title><xsl:value-of select="@title"/></title>
      </xsl:if>

  </xsl:template>

  <xsl:template match="p">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="ul">
    <list><xsl:apply-templates/></list>
  </xsl:template>

  <xsl:template match="ol">
    <enum><xsl:apply-templates/></enum>
  </xsl:template>

  <xsl:template match="li">
    <item><xsl:apply-templates/></item>
  </xsl:template>

  <xsl:template match="table">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="th">
    <xsl:apply-templates/>
  </xsl:template>

  <xsl:template match="tr">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="td">
    <co><xsl:apply-templates/></co>
  </xsl:template>

  <xsl:template match="pre">
    <source><xsl:apply-templates/></source>
  </xsl:template>

  <xsl:template match="em">
    <term><xsl:apply-templates/></term>
  </xsl:template>

  <xsl:template match="strong">
    <imp><xsl:apply-templates/></imp>
  </xsl:template>

  <xsl:template match="a">
    <link url="{@href}"><xsl:apply-templates/></link>
  </xsl:template>

</xsl:stylesheet>`
	XHTML_HEADER = "<xhtml>\n"
	XHTML_FOOTER = "\n</xhtml>"
)

func processXsl(xmlFile string, data map[string]string, article bool) []byte {
	xslFile, err := ioutil.TempFile("/tmp", "md2xsl-")
	if err != nil {
		panic(err)
	}
	stylesheet := STYLESHEET_ARTICLE
	if !article {
		stylesheet = STYLESHEET_BLOG
	}
	err = ioutil.WriteFile(xslFile.Name(), []byte(stylesheet), 0755)
	if err != nil {
		panic(err)
	}
	defer os.Remove(xslFile.Name())
	params := make([]string, 0, 2+3*len(data))
	for name, value := range data {
		params = append(params, "--stringparam")
		params = append(params, name)
		params = append(params, value)
	}
	params = append(params, xslFile.Name())
	params = append(params, xmlFile)
	command := exec.Command("xsltproc", params...)
	result, err := command.CombinedOutput()
	if err != nil {
	 println(result)
		panic(err)
	}
	return result
}

func markdown2xhtml(markdown string) []byte {
	mdFile, err := ioutil.TempFile("/tmp", "md2xsl-")
	if err != nil {
		panic(err)
	}
	defer os.Remove(mdFile.Name())
	ioutil.WriteFile(mdFile.Name(), []byte(markdown), 0x755)
	command := exec.Command("pandoc", mdFile.Name(), "-f", "markdown", "-t", "html")
	result, err := command.CombinedOutput()
	if err != nil {
	 println(result)
		panic(err)
	}
	return []byte(XHTML_HEADER + string(result) + XHTML_FOOTER)
}

func markdownData(text string) (map[string]string, string) {
	data := make(map[string]string)
	lines := strings.Split(text, "\n")
	var limit int
	for index, line := range lines {
		if strings.HasPrefix(line, "% ") && strings.Index(line, ":") >= 0 {
			name := strings.TrimSpace(line[2:strings.Index(line, ":")])
			value := strings.TrimSpace(line[strings.Index(line, ":")+1 : len(line)])
			data[name] = value
		} else {
			limit = index
			break
		}
	}
	return data, strings.Join(lines[limit:len(lines)], "\n")
}

func processFile(filename string, printXhtml bool, article bool) string {
	source, err := ioutil.ReadFile(filename)
	if err != nil {
		panic(err)
	}
	data, markdown := markdownData(string(source))
	xhtml := markdown2xhtml(markdown)
	if printXhtml {
		return string(xhtml)
	}
	xmlFile, err := ioutil.TempFile("/tmp", "md2xml-")
	if err != nil {
		panic(err)
	}
	defer os.Remove(xmlFile.Name())
	ioutil.WriteFile(xmlFile.Name(), xhtml, 0755)
	result := processXsl(xmlFile.Name(), data, article)
	return string(result)
}

func main() {
	xhtml := false
	article := false
	if len(os.Args) < 2 {
		fmt.Println(HELP)
		os.Exit(1)
	}
	for _, arg := range os.Args[1:] {
		if arg == "-h" || os.Args[1] == "--help" {
			fmt.Println(HELP)
			os.Exit(0)
		} else if arg == "-x" || arg == "--xhtml" {
			xhtml = true
		} else if arg == "-a" || arg == "--article" {
			article = true
		} else if arg == "-s" || arg == "--stylesheet" {
   if article {
       fmt.Println(STYLESHEET_ARTICLE)
   } else {
       fmt.Println(STYLESHEET_BLOG)
   }
		} else {
			fmt.Println(processFile(arg, xhtml, article))
		}
	}
}
<commit_msg>Added lang parameter<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"
)

const (
	HELP = `md2xml [-h] file.md
Transform a given Markdown file into XML.
-h        To print this help page.
-x        Print intermediate XHTML output.
-a        Output article (instead of blog entry).
-s        Print stylesheet used for transformation.
file.md   The markdown file to convert.
Note: this program calls xsltproc that must have been installed.`
	STYLESHEET_ARTICLE = `<?xml version="1.0" encoding="utf-8"?>

<xsl:stylesheet xmlns:xsl="http://www.w3.org/1999/XSL/Transform"
                version="1.0">

  <xsl:output method="xml" encoding="UTF-8"/>
  <xsl:param name="id">ID</xsl:param>
  <xsl:param name="date">DATE</xsl:param>
  <xsl:param name="title">TITLE</xsl:param>
  <xsl:param name="author">AUTHOR</xsl:param>
  <xsl:param name="email">EMAIL</xsl:param>
  <xsl:param name="lang">LANG</xsl:param>

  <xsl:template match="/xhtml">
    <xsl:text disable-output-escaping="yes">
    <!DOCTYPE article PUBLIC "-//CAFEBABE//DTD blog 1.0//EN"
                             "../dtd/article.dtd">
    </xsl:text>

      <xsl:attribute name="id"><xsl:value-of select="$id"/></xsl:attribute>
      <xsl:attribute name="date"><xsl:value-of select="$date"/></xsl:attribute>
      <xsl:attribute name="author"><xsl:value-of select="$author"/></xsl:attribute>
      <xsl:attribute name="email"><xsl:value-of select="$email"/></xsl:attribute>
      <xsl:attribute name="lang"><xsl:value-of select="$lang"/></xsl:attribute>
      <title><xsl:value-of select="$title"/></title>
      <text>
       <xsl:apply-templates/>
      </text>

  </xsl:template>

  <xsl:template match="h1">
    <sect level="1"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="h2">
    <sect level="2"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="h3">
    <sect level="3"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="h4">
    <sect level="4"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="h5">
    <sect level="5"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="h6">
    <sect level="6"><title><xsl:value-of select="."/></title></sect>
  </xsl:template>

  <xsl:template match="p[@class='caption']">
  </xsl:template>

  <xsl:template match="p[count(text())=0 and count(code)=1]">
    <source><xsl:apply-templates select="code"/></source>
  </xsl:template>

  <xsl:template match="p[count(text())=1 and count(img)=1]">
    <xsl:apply-templates select="img"/>
  </xsl:template>

  <xsl:template match="img">

      <xsl:if test="@title">
        <title><xsl:value-of select="@title"/></title>
      </xsl:if>

  </xsl:template>

  <xsl:template match="p">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="ul">
    <list><xsl:apply-templates/></list>
  </xsl:template>

  <xsl:template match="ol">
    <enum><xsl:apply-templates/></enum>
  </xsl:template>

  <xsl:template match="li">
    <item><xsl:apply-templates/></item>
  </xsl:template>

  <xsl:template match="table">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="th">
    <xsl:apply-templates/>
  </xsl:template>

  <xsl:template match="tr">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="td">
    <co><xsl:apply-templates/></co>
  </xsl:template>

  <xsl:template match="pre">
    <source><xsl:apply-templates/></source>
  </xsl:template>

  <xsl:template match="em">
    <term><xsl:apply-templates/></term>
  </xsl:template>

  <xsl:template match="strong">
    <imp><xsl:apply-templates/></imp>
  </xsl:template>

  <xsl:template match="a">
    <link url="{@href}"><xsl:apply-templates/></link>
  </xsl:template>

</xsl:stylesheet>`
	STYLESHEET_BLOG = `<?xml version="1.0" encoding="utf-8"?>

<xsl:stylesheet xmlns:xsl="http://www.w3.org/1999/XSL/Transform"
                version="1.0">

  <xsl:output method="xml" encoding="UTF-8"/>
  <xsl:param name="id">ID</xsl:param>
  <xsl:param name="date">DATE</xsl:param>
  <xsl:param name="title">TITLE</xsl:param>

  <xsl:template match="/xhtml">
    <xsl:text disable-output-escaping="yes">
    <!DOCTYPE blog PUBLIC "-//CAFEBABE//DTD blog 1.0//EN"
                          "../dtd/blog.dtd">
    </xsl:text>
    <blog>
      <xsl:attribute name="id"><xsl:value-of select="$id"/></xsl:attribute>
      <xsl:attribute name="date"><xsl:value-of select="$date"/></xsl:attribute>
      <title><xsl:value-of select="$title"/></title>
      <xsl:apply-templates/>
    </blog>
  </xsl:template>

  <xsl:template match="h1">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="h2">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="h3">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="h4">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="h5">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="h6">

<imp><xsl:value-of select="."/></imp>

  </xsl:template>

  <xsl:template match="p[@class='caption']">
  </xsl:template>

  <xsl:template match="p[count(text())=0 and count(code)=1]">
    <source><xsl:apply-templates select="code"/></source>
  </xsl:template>

  <xsl:template match="p[count(text())=1 and count(img)=1]">
    <xsl:apply-templates select="img"/>
  </xsl:template>

  <xsl:template match="img">

      <xsl:if test="@title">
        <title><xsl:value-of select="@title"/></title>
      </xsl:if>

  </xsl:template>

  <xsl:template match="p">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="ul">
    <list><xsl:apply-templates/></list>
  </xsl:template>

  <xsl:template match="ol">
    <enum><xsl:apply-templates/></enum>
  </xsl:template>

  <xsl:template match="li">
    <item><xsl:apply-templates/></item>
  </xsl:template>

  <xsl:template match="table">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="th">
    <xsl:apply-templates/>
  </xsl:template>

  <xsl:template match="tr">

<xsl:apply-templates/>

  </xsl:template>

  <xsl:template match="td">
    <co><xsl:apply-templates/></co>
  </xsl:template>

  <xsl:template match="pre">
    <source><xsl:apply-templates/></source>
  </xsl:template>

  <xsl:template match="em">
    <term><xsl:apply-templates/></term>
  </xsl:template>

  <xsl:template match="strong">
    <imp><xsl:apply-templates/></imp>
  </xsl:template>

  <xsl:template match="a">
    <link url="{@href}"><xsl:apply-templates/></link>
  </xsl:template>

</xsl:stylesheet>`
	XHTML_HEADER = "<xhtml>\n"
	XHTML_FOOTER = "\n</xhtml>"
)

func processXsl(xmlFile string, data map[string]string, article bool) []byte {
	xslFile, err := ioutil.TempFile("/tmp", "md2xsl-")
	if err != nil {
		panic(err)
	}
	stylesheet := STYLESHEET_ARTICLE
	if !article {
		stylesheet = STYLESHEET_BLOG
	}
	err = ioutil.WriteFile(xslFile.Name(), []byte(stylesheet), 0755)
	if err != nil {
		panic(err)
	}
	defer os.Remove(xslFile.Name())
	params := make([]string, 0, 2+3*len(data))
	for name, value := range data {
		params = append(params, "--stringparam")
		params = append(params, name)
		params = append(params, value)
	}
	params = append(params, xslFile.Name())
	params = append(params, xmlFile)
	command := exec.Command("xsltproc", params...)
	result, err := command.CombinedOutput()
	if err != nil {
        println(result)
		panic(err)
	}
	return result
}

func markdown2xhtml(markdown string) []byte {
	mdFile, err := ioutil.TempFile("/tmp", "md2xsl-")
	if err != nil {
		panic(err)
	}
	defer os.Remove(mdFile.Name())
	ioutil.WriteFile(mdFile.Name(), []byte(markdown), 0x755)
	command := exec.Command("pandoc", mdFile.Name(), "-f", "markdown", "-t", "html")
	result, err := command.CombinedOutput()
	if err != nil {
	 println(result)
		panic(err)
	}
	return []byte(XHTML_HEADER + string(result) + XHTML_FOOTER)
}

func markdownData(text string) (map[string]string, string) {
	data := make(map[string]string)
	lines := strings.Split(text, "\n")
	var limit int
	for index, line := range lines {
		if strings.HasPrefix(line, "% ") && strings.Index(line, ":") >= 0 {
			name := strings.TrimSpace(line[2:strings.Index(line, ":")])
			value := strings.TrimSpace(line[strings.Index(line, ":")+1 : len(line)])
			data[name] = value
		} else {
			limit = index
			break
		}
	}
	return data, strings.Join(lines[limit:len(lines)], "\n")
}

func processFile(filename string, printXhtml bool, article bool) string {
	source, err := ioutil.ReadFile(filename)
	if err != nil {
		panic(err)
	}
	data, markdown := markdownData(string(source))
	xhtml := markdown2xhtml(markdown)
	if printXhtml {
		return string(xhtml)
	}
	xmlFile, err := ioutil.TempFile("/tmp", "md2xml-")
	if err != nil {
		panic(err)
	}
	defer os.Remove(xmlFile.Name())
	ioutil.WriteFile(xmlFile.Name(), xhtml, 0755)
	result := processXsl(xmlFile.Name(), data, article)
	return string(result)
}

func main() {
	xhtml := false
	article := false
	if len(os.Args) < 2 {
		fmt.Println(HELP)
		os.Exit(1)
	}
	for _, arg := range os.Args[1:] {
		if arg == "-h" || os.Args[1] == "--help" {
			fmt.Println(HELP)
			os.Exit(0)
		} else if arg == "-x" || arg == "--xhtml" {
			xhtml = true
		} else if arg == "-a" || arg == "--article" {
			article = true
		} else if arg == "-s" || arg == "--stylesheet" {
   if article {
       fmt.Println(STYLESHEET_ARTICLE)
   } else {
       fmt.Println(STYLESHEET_BLOG)
   }
		} else {
			fmt.Println(processFile(arg, xhtml, article))
		}
	}
}
<|endoftext|>