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

// Package whirlpool implements the ISO/IEC 10118-3:2004 whirlpool
// cryptographic hash. Whirlpool is defined in
// http://www.larc.usp.br/~pbarreto/WhirlpoolPage.html

package whirlpool

import (
	"encoding/binary"
	"hash"
)

type whirlpool struct {
	bitLength  [lengthBytes]byte       // Number of hashed bits.
	buffer     [wblockBytes]byte       // Buffer of data to be hashed.
	bufferBits int                     // Current number of bits on the buffer.
	bufferPos  int                     // Current byte location on buffer.
	hash       [digestBytes / 8]uint64 // Hash state.
}

func New() hash.Hash {
	return new(whirlpool)
}

func (w *whirlpool) Reset() {
	// Cleanup the buffer.
	w.buffer = [wblockBytes]byte{}
	w.bufferBits = 0
	w.bufferPos = 0

	// Cleanup the digest.
	w.hash = [digestBytes / 8]uint64{}

	// Clean up the number of hashed bits.
	w.bitLength = [lengthBytes]byte{}
}

func (w *whirlpool) Size() int {
	return digestBytes
}

func (w *whirlpool) BlockSize() int {
	return wblockBytes
}

func (w *whirlpool) transform() {
	var (
		K     [8]uint64 // Round key.
		block [8]uint64 // μ(buffer).
		state [8]uint64 // Cipher state.
		L     [8]uint64
	)

	// Map the buffer to a block.
	for i := 0; i < 8; i++ {
		b := 8 * i
		block[i] = binary.BigEndian.Uint64(w.buffer[b:])
	}

	// Compute & apply K^0 to the cipher state.
	for i := 0; i < 8; i++ {
		K[i] = w.hash[i]
		state[i] = block[i] ^ K[i]
	}

	// Iterate over all the rounds.
	for r := 1; r <= rounds; r++ {
		// Compute K^rounds from K^(rounds-1).
		for i := 0; i < 8; i++ {
			L[i] = C0[byte(K[i%8]>>56)] ^
				C1[byte(K[(i+7)%8]>>48)] ^
				C2[byte(K[(i+6)%8]>>40)] ^
				C3[byte(K[(i+5)%8]>>32)] ^
				C4[byte(K[(i+4)%8]>>24)] ^
				C5[byte(K[(i+3)%8]>>16)] ^
				C6[byte(K[(i+2)%8]>>8)] ^
				C7[byte(K[(i+1)%8])]
		}
		L[0] ^= rc[r]

		for i := 0; i < 8; i++ {
			K[i] = L[i]
		}

		// Apply r-th round transformation.
		for i := 0; i < 8; i++ {
			L[i] = C0[byte(state[i%8]>>56)] ^
				C1[byte(state[(i+7)%8]>>48)] ^
				C2[byte(state[(i+6)%8]>>40)] ^
				C3[byte(state[(i+5)%8]>>32)] ^
				C4[byte(state[(i+4)%8]>>24)] ^
				C5[byte(state[(i+3)%8]>>16)] ^
				C6[byte(state[(i+2)%8]>>8)] ^
				C7[byte(state[(i+1)%8])] ^
				K[i%8]
		}

		for i := 0; i < 8; i++ {
			state[i] = L[i]
		}
	}

	// Apply the Miyaguchi-Preneel compression function.
	for i := 0; i < 8; i++ {
		w.hash[i] ^= state[i] ^ block[i]
	}
}

func (w *whirlpool) Write(source []byte) (int, error) {
	var (
		sourcePos  int                                            // Index of the leftmost source.
		nn         int    = len(source)                           // Num of bytes to process.
		sourceBits uint64 = uint64(nn * 8)                        // Num of bits to process.
		sourceGap  uint   = uint((8 - (int(sourceBits & 7))) & 7) // Space on source[sourcePos].
		bufferRem  uint   = uint(w.bufferBits & 7)                // Occupied bits on buffer[bufferPos].
		b          uint32                                         // Current byte.
	)

	// Tally the length of the data added.
	for i, carry, value := 31, uint32(0), uint64(sourceBits); i >= 0 && (carry != 0 || value != 0); i-- {
		carry += uint32(w.bitLength[i]) + (uint32(value & 0xff))
		w.bitLength[i] = byte(carry)
		carry >>= 8
		value >>= 8
	}

	// Process data in chunks of 8 bits.
	for sourceBits > 8 {
		// Take a byte form the source.
		b = uint32(((source[sourcePos] << sourceGap) & 0xff) |
			((source[sourcePos+1] & 0xff) >> (8 - sourceGap)))

		// Process this byte.
		w.buffer[w.bufferPos] |= uint8(b >> bufferRem)
		w.bufferPos++
		w.bufferBits += int(8 - bufferRem)

		if w.bufferBits == digestBits {
			// Process this block.
			w.transform()
			// Reset the buffer.
			w.bufferBits = 0
			w.bufferPos = 0
		}
		w.buffer[w.bufferPos] = byte(b << (8 - bufferRem))
		w.bufferBits += int(bufferRem)

		// Proceed to remaining data.
		sourceBits -= 8
		sourcePos++
	}

	// 0 <= sourceBits <= 8; All data leftover is in source[sourcePos].
	if sourceBits > 0 {
		b = uint32((source[sourcePos] << sourceGap) & 0xff) // bits are left-justified

		// Process the remaining bits.
		w.buffer[w.bufferPos] |= byte(b) >> bufferRem
	} else {
		b = 0
	}

	if uint64(bufferRem)+sourceBits < 8 {
		// The remaining data fits on the buffer[bufferPos].
		w.bufferBits += int(sourceBits)
	} else {
		// The buffer[bufferPos] is full.
		w.bufferPos++
		w.bufferBits += 8 - int(bufferRem) // bufferBits = 8*bufferPos
		sourceBits -= uint64(8 - bufferRem)

		// Now, 0 <= sourceBits <= 8; all data leftover is in source[sourcePos].
		if w.bufferBits == digestBits {
			// Process this data block.
			w.transform()
			// Reset buffer.
			w.bufferBits = 0
			w.bufferPos = 0
		}
		w.buffer[w.bufferPos] = byte(b << (8 - bufferRem))
		w.bufferBits += int(sourceBits)
	}
	return nn, nil
}

func (w *whirlpool) Sum(in []byte) []byte {
	// Copy the whirlpool so that the caller can keep summing.
	n := *w

	// Append a 1-bit.
	n.buffer[n.bufferPos] |= 0x80 >> (uint(n.bufferBits) & 7)
	n.bufferPos++

	// The remaining bits should be 0. Pad with 0s to be complete.
	if n.bufferPos > wblockBytes-lengthBytes {
		if n.bufferPos < wblockBytes {
			for i := 0; i < wblockBytes-n.bufferPos; i++ {
				n.buffer[n.bufferPos+i] = 0
			}
		}
		// Process this data block.
		n.transform()
		// Reset the buffer.
		n.bufferPos = 0
	}

	if n.bufferPos < wblockBytes-lengthBytes {
		for i := 0; i < (wblockBytes-lengthBytes)-n.bufferPos; i++ {
			n.buffer[n.bufferPos+i] = 0
		}
	}
	n.bufferPos = wblockBytes - lengthBytes

	// Append the bit length of the hashed data.
	for i := 0; i < lengthBytes; i++ {
		n.buffer[n.bufferPos+i] = n.bitLength[i]
	}

	// Process this data block.
	n.transform()

	// Return the final digest as []byte.
	var digest [digestBytes]byte
	for i := 0; i < digestBytes/8; i++ {
		digest[i*8] = byte(n.hash[i] >> 56)
		digest[i*8+1] = byte(n.hash[i] >> 48)
		digest[i*8+2] = byte(n.hash[i] >> 40)
		digest[i*8+3] = byte(n.hash[i] >> 32)
		digest[i*8+4] = byte(n.hash[i] >> 24)
		digest[i*8+5] = byte(n.hash[i] >> 16)
		digest[i*8+6] = byte(n.hash[i] >> 8)
		digest[i*8+7] = byte(n.hash[i])
	}

	return append(in, digest[:digestBytes]...)
}
<commit_msg>added comments for better doc<commit_after>// Copyright 2012 Jimmy Zelinskie. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package whirlpool implements the ISO/IEC 10118-3:2004 whirlpool
// cryptographic hash. Whirlpool is defined in
// http://www.larc.usp.br/~pbarreto/WhirlpoolPage.html

package whirlpool

import (
	"encoding/binary"
	"hash"
)

// whirlpool represents the partial evaluation of a checksum.
type whirlpool struct {
	bitLength  [lengthBytes]byte       // Number of hashed bits.
	buffer     [wblockBytes]byte       // Buffer of data to be hashed.
	bufferBits int                     // Current number of bits on the buffer.
	bufferPos  int                     // Current byte location on buffer.
	hash       [digestBytes / 8]uint64 // Hash state.
}

// New returns a new hash.Hash computing the whirlpool checksum.
func New() hash.Hash {
	return new(whirlpool)
}

func (w *whirlpool) Reset() {
	// Cleanup the buffer.
	w.buffer = [wblockBytes]byte{}
	w.bufferBits = 0
	w.bufferPos = 0

	// Cleanup the digest.
	w.hash = [digestBytes / 8]uint64{}

	// Clean up the number of hashed bits.
	w.bitLength = [lengthBytes]byte{}
}

func (w *whirlpool) Size() int {
	return digestBytes
}

func (w *whirlpool) BlockSize() int {
	return wblockBytes
}

func (w *whirlpool) transform() {
	var (
		K     [8]uint64 // Round key.
		block [8]uint64 // μ(buffer).
		state [8]uint64 // Cipher state.
		L     [8]uint64
	)

	// Map the buffer to a block.
	for i := 0; i < 8; i++ {
		b := 8 * i
		block[i] = binary.BigEndian.Uint64(w.buffer[b:])
	}

	// Compute & apply K^0 to the cipher state.
	for i := 0; i < 8; i++ {
		K[i] = w.hash[i]
		state[i] = block[i] ^ K[i]
	}

	// Iterate over all the rounds.
	for r := 1; r <= rounds; r++ {
		// Compute K^rounds from K^(rounds-1).
		for i := 0; i < 8; i++ {
			L[i] = C0[byte(K[i%8]>>56)] ^
				C1[byte(K[(i+7)%8]>>48)] ^
				C2[byte(K[(i+6)%8]>>40)] ^
				C3[byte(K[(i+5)%8]>>32)] ^
				C4[byte(K[(i+4)%8]>>24)] ^
				C5[byte(K[(i+3)%8]>>16)] ^
				C6[byte(K[(i+2)%8]>>8)] ^
				C7[byte(K[(i+1)%8])]
		}
		L[0] ^= rc[r]

		for i := 0; i < 8; i++ {
			K[i] = L[i]
		}

		// Apply r-th round transformation.
		for i := 0; i < 8; i++ {
			L[i] = C0[byte(state[i%8]>>56)] ^
				C1[byte(state[(i+7)%8]>>48)] ^
				C2[byte(state[(i+6)%8]>>40)] ^
				C3[byte(state[(i+5)%8]>>32)] ^
				C4[byte(state[(i+4)%8]>>24)] ^
				C5[byte(state[(i+3)%8]>>16)] ^
				C6[byte(state[(i+2)%8]>>8)] ^
				C7[byte(state[(i+1)%8])] ^
				K[i%8]
		}

		for i := 0; i < 8; i++ {
			state[i] = L[i]
		}
	}

	// Apply the Miyaguchi-Preneel compression function.
	for i := 0; i < 8; i++ {
		w.hash[i] ^= state[i] ^ block[i]
	}
}

func (w *whirlpool) Write(source []byte) (int, error) {
	var (
		sourcePos  int                                            // Index of the leftmost source.
		nn         int    = len(source)                           // Num of bytes to process.
		sourceBits uint64 = uint64(nn * 8)                        // Num of bits to process.
		sourceGap  uint   = uint((8 - (int(sourceBits & 7))) & 7) // Space on source[sourcePos].
		bufferRem  uint   = uint(w.bufferBits & 7)                // Occupied bits on buffer[bufferPos].
		b          uint32                                         // Current byte.
	)

	// Tally the length of the data added.
	for i, carry, value := 31, uint32(0), uint64(sourceBits); i >= 0 && (carry != 0 || value != 0); i-- {
		carry += uint32(w.bitLength[i]) + (uint32(value & 0xff))
		w.bitLength[i] = byte(carry)
		carry >>= 8
		value >>= 8
	}

	// Process data in chunks of 8 bits.
	for sourceBits > 8 {
		// Take a byte form the source.
		b = uint32(((source[sourcePos] << sourceGap) & 0xff) |
			((source[sourcePos+1] & 0xff) >> (8 - sourceGap)))

		// Process this byte.
		w.buffer[w.bufferPos] |= uint8(b >> bufferRem)
		w.bufferPos++
		w.bufferBits += int(8 - bufferRem)

		if w.bufferBits == digestBits {
			// Process this block.
			w.transform()
			// Reset the buffer.
			w.bufferBits = 0
			w.bufferPos = 0
		}
		w.buffer[w.bufferPos] = byte(b << (8 - bufferRem))
		w.bufferBits += int(bufferRem)

		// Proceed to remaining data.
		sourceBits -= 8
		sourcePos++
	}

	// 0 <= sourceBits <= 8; All data leftover is in source[sourcePos].
	if sourceBits > 0 {
		b = uint32((source[sourcePos] << sourceGap) & 0xff) // The bits are left-justified.

		// Process the remaining bits.
		w.buffer[w.bufferPos] |= byte(b) >> bufferRem
	} else {
		b = 0
	}

	if uint64(bufferRem)+sourceBits < 8 {
		// The remaining data fits on the buffer[bufferPos].
		w.bufferBits += int(sourceBits)
	} else {
		// The buffer[bufferPos] is full.
		w.bufferPos++
		w.bufferBits += 8 - int(bufferRem) // bufferBits = 8*bufferPos
		sourceBits -= uint64(8 - bufferRem)

		// Now, 0 <= sourceBits <= 8; all data leftover is in source[sourcePos].
		if w.bufferBits == digestBits {
			// Process this data block.
			w.transform()
			// Reset buffer.
			w.bufferBits = 0
			w.bufferPos = 0
		}
		w.buffer[w.bufferPos] = byte(b << (8 - bufferRem))
		w.bufferBits += int(sourceBits)
	}
	return nn, nil
}

func (w *whirlpool) Sum(in []byte) []byte {
	// Copy the whirlpool so that the caller can keep summing.
	n := *w

	// Append a 1-bit.
	n.buffer[n.bufferPos] |= 0x80 >> (uint(n.bufferBits) & 7)
	n.bufferPos++

	// The remaining bits should be 0. Pad with 0s to be complete.
	if n.bufferPos > wblockBytes-lengthBytes {
		if n.bufferPos < wblockBytes {
			for i := 0; i < wblockBytes-n.bufferPos; i++ {
				n.buffer[n.bufferPos+i] = 0
			}
		}
		// Process this data block.
		n.transform()
		// Reset the buffer.
		n.bufferPos = 0
	}

	if n.bufferPos < wblockBytes-lengthBytes {
		for i := 0; i < (wblockBytes-lengthBytes)-n.bufferPos; i++ {
			n.buffer[n.bufferPos+i] = 0
		}
	}
	n.bufferPos = wblockBytes - lengthBytes

	// Append the bit length of the hashed data.
	for i := 0; i < lengthBytes; i++ {
		n.buffer[n.bufferPos+i] = n.bitLength[i]
	}

	// Process this data block.
	n.transform()

	// Return the final digest as []byte.
	var digest [digestBytes]byte
	for i := 0; i < digestBytes/8; i++ {
		digest[i*8] = byte(n.hash[i] >> 56)
		digest[i*8+1] = byte(n.hash[i] >> 48)
		digest[i*8+2] = byte(n.hash[i] >> 40)
		digest[i*8+3] = byte(n.hash[i] >> 32)
		digest[i*8+4] = byte(n.hash[i] >> 24)
		digest[i*8+5] = byte(n.hash[i] >> 16)
		digest[i*8+6] = byte(n.hash[i] >> 8)
		digest[i*8+7] = byte(n.hash[i])
	}

	return append(in, digest[:digestBytes]...)
}
<|endoftext|>


<commit_before>/*
Copyright 2022 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package prometheusextension

import (
	"math"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
)

func TestTimingHistogramNonMonotonicBuckets(t *testing.T) {
	testCases := map[string][]float64{
		"not strictly monotonic":  {1, 2, 2, 3},
		"not monotonic at all":    {1, 2, 4, 3, 5},
		"have +Inf in the middle": {1, 2, math.Inf(+1), 3},
	}
	for name, buckets := range testCases {
		_, err := NewTimingHistogram(TimingHistogramOpts{
			Name:    "test_histogram",
			Help:    "helpless",
			Buckets: buckets,
		})
		if err == nil {
			t.Errorf("Buckets %v are %s but NewHTimingistogram did not complain.", buckets, name)
		}
	}
}

func exerciseTimingHistogramAndCollector(th GaugeOps, t0 time.Time, clk *unsyncFakeClock, collect func(chan<- prometheus.Metric), expectCollection ...GaugeOps) func(t *testing.T) {
	return func(t *testing.T) {
		exerciseTimingHistogramData(t, th, t0, clk)
		exerciseTimingHistogramCollector(t, collect, expectCollection)
	}
}

func exerciseTimingHistogramCollector(t *testing.T, collect func(chan<- prometheus.Metric), expectCollection []GaugeOps) {
	remainingCollection := expectCollection
	metch := make(chan prometheus.Metric)
	go func() {
		collect(metch)
		close(metch)
	}()
	for collected := range metch {
		collectedGO := collected.(GaugeOps)
		newRem, found := findAndRemove(remainingCollection, collectedGO)
		if !found {
			t.Errorf("Collected unexpected value %#+v", collected)
		}
		remainingCollection = newRem
	}
	if len(remainingCollection) > 0 {
		t.Errorf("Collection omitted %#+v", remainingCollection)
	}
}

var thTestBuckets = []float64{0, 0.5, 1}
var thTestV0 float64 = 0.25

// exerciseTimingHistogramData takes the given histogram through the following points in (time,value) space.
// t0 is the clock time of the histogram's construction
// value=v0 for t0 <= t <= t1 where v0 = 0.25 and t1 = t0 + 1 ns
// value=v1 for t1 <= t <= t2 where v1 = 0.75 and t2 = t1 + 1 microsecond
// value=v2 for t2 <= t <= t3 where v2 = 1.25 and t3 = t2 + 1 millisecond
// value=v3 for t3 <= t <= t4 where v3 = 0.65 and t4 = t3 + 1 second
func exerciseTimingHistogramData(t *testing.T, th GaugeOps, t0 time.Time, clk *unsyncFakeClock) {
	t1 := t0.Add(time.Nanosecond)
	v0 := thTestV0
	var v1 float64 = 0.75
	clk.SetTime(t1)
	th.Set(v1)
	t2 := t1.Add(time.Microsecond)
	var d2 float64 = 0.5
	v2 := v1 + d2
	clk.SetTime(t2)
	th.Add(d2)
	t3 := t2
	for i := 0; i < 1000000; i++ {
		t3 = t3.Add(time.Nanosecond)
		clk.SetTime(t3)
		th.Set(v2)
	}
	var d3 float64 = -0.6
	v3 := v2 + d3
	th.Add(d3)
	t4 := t3.Add(time.Second)
	clk.SetTime(t4)

	metric := &dto.Metric{}
	writer := th.(prometheus.Metric)
	err := writer.Write(metric)
	if err != nil {
		t.Error(err)
	}
	wroteHist := metric.Histogram
	if want, got := uint64(t4.Sub(t0)), wroteHist.GetSampleCount(); want != got {
		t.Errorf("Wanted %v but got %v", want, got)
	}
	if want, got := tDiff(t1, t0)*v0+tDiff(t2, t1)*v1+tDiff(t3, t2)*v2+tDiff(t4, t3)*v3, wroteHist.GetSampleSum(); want != got {
		t.Errorf("Wanted %v but got %v", want, got)
	}
	wroteBuckets := wroteHist.GetBucket()
	if len(wroteBuckets) != len(thTestBuckets) {
		t.Errorf("Got buckets %#+v", wroteBuckets)
	}
	expectedCounts := []time.Duration{0, t1.Sub(t0), t2.Sub(t0) + t4.Sub(t3)}
	for idx, ub := range thTestBuckets {
		if want, got := uint64(expectedCounts[idx]), wroteBuckets[idx].GetCumulativeCount(); want != got {
			t.Errorf("In bucket %d, wanted %v but got %v", idx, want, got)
		}
		if want, got := ub, wroteBuckets[idx].GetUpperBound(); want != got {
			t.Errorf("In bucket %d, wanted %v but got %v", idx, want, got)
		}
	}
}

// tDiff returns a time difference as float
func tDiff(hi, lo time.Time) float64 { return float64(hi.Sub(lo)) }

func findAndRemove(metrics []GaugeOps, seek GaugeOps) ([]GaugeOps, bool) {
	for idx, metric := range metrics {
		if metric == seek {
			return append(append([]GaugeOps{}, metrics[:idx]...), metrics[idx+1:]...), true
		}
	}
	return metrics, false
}

func TestTimeIntegrationDirect(t *testing.T) {
	t0 := time.Now()
	clk := &unsyncFakeClock{t0}
	th, err := NewTestableTimingHistogram(clk.Now, TimingHistogramOpts{
		Name:         "TestTimeIntegration",
		Help:         "helpless",
		Buckets:      thTestBuckets,
		InitialValue: thTestV0,
	})
	if err != nil {
		t.Error(err)
		return
	}
	t.Run("non-vec", exerciseTimingHistogramAndCollector(th, t0, clk, th.Collect, th))
}

func TestTimingHistogramVec(t *testing.T) {
	t0 := time.Now()
	clk := &unsyncFakeClock{t0}
	vec := NewTestableTimingHistogramVec(clk.Now, TimingHistogramOpts{
		Name:         "TestTimeIntegration",
		Help:         "helpless",
		Buckets:      thTestBuckets,
		InitialValue: thTestV0,
	}, "k1", "k2")
	th1 := vec.With(prometheus.Labels{"k1": "a", "k2": "x"})
	th1b := vec.WithLabelValues("a", "x")
	if th1 != th1b {
		t.Errorf("Vector not functional")
	}
	t.Run("th1", exerciseTimingHistogramAndCollector(th1, t0, clk, vec.Collect, th1))
	t0 = clk.Now()
	th2 := vec.WithLabelValues("a", "y")
	if th1 == th2 {
		t.Errorf("Vector does not distinguish label values")
	}
	t.Run("th2", exerciseTimingHistogramAndCollector(th2, t0, clk, vec.Collect, th1, th2))
	t0 = clk.Now()
	th3 := vec.WithLabelValues("b", "y")
	if th1 == th3 || th2 == th3 {
		t.Errorf("Vector does not distinguish label values")
	}
	t.Run("th2", exerciseTimingHistogramAndCollector(th3, t0, clk, vec.Collect, th1, th2, th3))
}

type unsyncFakeClock struct {
	now time.Time
}

func (ufc *unsyncFakeClock) Now() time.Time {
	return ufc.now
}

func (ufc *unsyncFakeClock) SetTime(now time.Time) {
	ufc.now = now
}

func BenchmarkTimingHistogramDirect(b *testing.B) {
	b.StopTimer()
	now := time.Now()
	clk := &unsyncFakeClock{now: now}
	hist, err := NewTestableTimingHistogram(clk.Now, TimingHistogramOpts{
		Namespace: "testns",
		Subsystem: "testsubsys",
		Name:      "testhist",
		Help:      "Me",
		Buckets:   []float64{1, 2, 4, 8, 16},
	})
	if err != nil {
		b.Error(err)
	}
	var x int
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		clk.now = clk.now.Add(time.Duration(31-x) * time.Microsecond)
		hist.Set(float64(x))
		x = (x + i) % 23
	}
}
<commit_msg>Simplify and extend the benchmarks of timing histograms<commit_after>/*
Copyright 2022 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package prometheusextension

import (
	"math"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	dto "github.com/prometheus/client_model/go"
)

func TestTimingHistogramNonMonotonicBuckets(t *testing.T) {
	testCases := map[string][]float64{
		"not strictly monotonic":  {1, 2, 2, 3},
		"not monotonic at all":    {1, 2, 4, 3, 5},
		"have +Inf in the middle": {1, 2, math.Inf(+1), 3},
	}
	for name, buckets := range testCases {
		_, err := NewTimingHistogram(TimingHistogramOpts{
			Name:    "test_histogram",
			Help:    "helpless",
			Buckets: buckets,
		})
		if err == nil {
			t.Errorf("Buckets %v are %s but NewHTimingistogram did not complain.", buckets, name)
		}
	}
}

func exerciseTimingHistogramAndCollector(th GaugeOps, t0 time.Time, clk *unsyncFakeClock, collect func(chan<- prometheus.Metric), expectCollection ...GaugeOps) func(t *testing.T) {
	return func(t *testing.T) {
		exerciseTimingHistogramData(t, th, t0, clk)
		exerciseTimingHistogramCollector(t, collect, expectCollection)
	}
}

func exerciseTimingHistogramCollector(t *testing.T, collect func(chan<- prometheus.Metric), expectCollection []GaugeOps) {
	remainingCollection := expectCollection
	metch := make(chan prometheus.Metric)
	go func() {
		collect(metch)
		close(metch)
	}()
	for collected := range metch {
		collectedGO := collected.(GaugeOps)
		newRem, found := findAndRemove(remainingCollection, collectedGO)
		if !found {
			t.Errorf("Collected unexpected value %#+v", collected)
		}
		remainingCollection = newRem
	}
	if len(remainingCollection) > 0 {
		t.Errorf("Collection omitted %#+v", remainingCollection)
	}
}

var thTestBuckets = []float64{0, 0.5, 1}
var thTestV0 float64 = 0.25

// exerciseTimingHistogramData takes the given histogram through the following points in (time,value) space.
// t0 is the clock time of the histogram's construction
// value=v0 for t0 <= t <= t1 where v0 = 0.25 and t1 = t0 + 1 ns
// value=v1 for t1 <= t <= t2 where v1 = 0.75 and t2 = t1 + 1 microsecond
// value=v2 for t2 <= t <= t3 where v2 = 1.25 and t3 = t2 + 1 millisecond
// value=v3 for t3 <= t <= t4 where v3 = 0.65 and t4 = t3 + 1 second
func exerciseTimingHistogramData(t *testing.T, th GaugeOps, t0 time.Time, clk *unsyncFakeClock) {
	t1 := t0.Add(time.Nanosecond)
	v0 := thTestV0
	var v1 float64 = 0.75
	clk.SetTime(t1)
	th.Set(v1)
	t2 := t1.Add(time.Microsecond)
	var d2 float64 = 0.5
	v2 := v1 + d2
	clk.SetTime(t2)
	th.Add(d2)
	t3 := t2
	for i := 0; i < 1000000; i++ {
		t3 = t3.Add(time.Nanosecond)
		clk.SetTime(t3)
		th.Set(v2)
	}
	var d3 float64 = -0.6
	v3 := v2 + d3
	th.Add(d3)
	t4 := t3.Add(time.Second)
	clk.SetTime(t4)

	metric := &dto.Metric{}
	writer := th.(prometheus.Metric)
	err := writer.Write(metric)
	if err != nil {
		t.Error(err)
	}
	wroteHist := metric.Histogram
	if want, got := uint64(t4.Sub(t0)), wroteHist.GetSampleCount(); want != got {
		t.Errorf("Wanted %v but got %v", want, got)
	}
	if want, got := tDiff(t1, t0)*v0+tDiff(t2, t1)*v1+tDiff(t3, t2)*v2+tDiff(t4, t3)*v3, wroteHist.GetSampleSum(); want != got {
		t.Errorf("Wanted %v but got %v", want, got)
	}
	wroteBuckets := wroteHist.GetBucket()
	if len(wroteBuckets) != len(thTestBuckets) {
		t.Errorf("Got buckets %#+v", wroteBuckets)
	}
	expectedCounts := []time.Duration{0, t1.Sub(t0), t2.Sub(t0) + t4.Sub(t3)}
	for idx, ub := range thTestBuckets {
		if want, got := uint64(expectedCounts[idx]), wroteBuckets[idx].GetCumulativeCount(); want != got {
			t.Errorf("In bucket %d, wanted %v but got %v", idx, want, got)
		}
		if want, got := ub, wroteBuckets[idx].GetUpperBound(); want != got {
			t.Errorf("In bucket %d, wanted %v but got %v", idx, want, got)
		}
	}
}

// tDiff returns a time difference as float
func tDiff(hi, lo time.Time) float64 { return float64(hi.Sub(lo)) }

func findAndRemove(metrics []GaugeOps, seek GaugeOps) ([]GaugeOps, bool) {
	for idx, metric := range metrics {
		if metric == seek {
			return append(append([]GaugeOps{}, metrics[:idx]...), metrics[idx+1:]...), true
		}
	}
	return metrics, false
}

func TestTimeIntegrationDirect(t *testing.T) {
	t0 := time.Now()
	clk := &unsyncFakeClock{t0}
	th, err := NewTestableTimingHistogram(clk.Now, TimingHistogramOpts{
		Name:         "TestTimeIntegration",
		Help:         "helpless",
		Buckets:      thTestBuckets,
		InitialValue: thTestV0,
	})
	if err != nil {
		t.Error(err)
		return
	}
	t.Run("non-vec", exerciseTimingHistogramAndCollector(th, t0, clk, th.Collect, th))
}

func TestTimingHistogramVec(t *testing.T) {
	t0 := time.Now()
	clk := &unsyncFakeClock{t0}
	vec := NewTestableTimingHistogramVec(clk.Now, TimingHistogramOpts{
		Name:         "TestTimeIntegration",
		Help:         "helpless",
		Buckets:      thTestBuckets,
		InitialValue: thTestV0,
	}, "k1", "k2")
	th1 := vec.With(prometheus.Labels{"k1": "a", "k2": "x"})
	th1b := vec.WithLabelValues("a", "x")
	if th1 != th1b {
		t.Errorf("Vector not functional")
	}
	t.Run("th1", exerciseTimingHistogramAndCollector(th1, t0, clk, vec.Collect, th1))
	t0 = clk.Now()
	th2 := vec.WithLabelValues("a", "y")
	if th1 == th2 {
		t.Errorf("Vector does not distinguish label values")
	}
	t.Run("th2", exerciseTimingHistogramAndCollector(th2, t0, clk, vec.Collect, th1, th2))
	t0 = clk.Now()
	th3 := vec.WithLabelValues("b", "y")
	if th1 == th3 || th2 == th3 {
		t.Errorf("Vector does not distinguish label values")
	}
	t.Run("th2", exerciseTimingHistogramAndCollector(th3, t0, clk, vec.Collect, th1, th2, th3))
}

type unsyncFakeClock struct {
	now time.Time
}

func (ufc *unsyncFakeClock) Now() time.Time {
	return ufc.now
}

func (ufc *unsyncFakeClock) SetTime(now time.Time) {
	ufc.now = now
}

func BenchmarkTimingHistogramDirect(b *testing.B) {
	b.StopTimer()
	now := time.Now()
	hist, err := NewTestableTimingHistogram(func() time.Time { return now }, TimingHistogramOpts{
		Namespace: "testns",
		Subsystem: "testsubsys",
		Name:      "testhist",
		Help:      "Me",
		Buckets:   []float64{1, 2, 4, 8, 16},
	})
	if err != nil {
		b.Error(err)
	}
	var x int
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		now = now.Add(time.Duration(31-x) * time.Microsecond)
		hist.Set(float64(x))
		x = (x + i) % 23
	}
}
func BenchmarkTimingHistogramVecEltCached(b *testing.B) {
	b.StopTimer()
	now := time.Now()
	vec := NewTestableTimingHistogramVec(func() time.Time { return now }, TimingHistogramOpts{
		Namespace: "testns",
		Subsystem: "testsubsys",
		Name:      "testhist",
		Help:      "Me",
		Buckets:   []float64{1, 2, 4, 8, 16},
	},
		"label1", "label2")
	hist := vec.WithLabelValues("val1", "val2")
	var x int
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		now = now.Add(time.Duration(31-x) * time.Microsecond)
		hist.Set(float64(x))
		x = (x + i) % 23
	}
}

func BenchmarkTimingHistogramVecEltFetched(b *testing.B) {
	b.StopTimer()
	now := time.Now()
	vec := NewTestableTimingHistogramVec(func() time.Time { return now }, TimingHistogramOpts{
		Namespace: "testns",
		Subsystem: "testsubsys",
		Name:      "testhist",
		Help:      "Me",
		Buckets:   []float64{1, 2, 4, 8, 16},
	},
		"label1", "label2")
	var x int
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		now = now.Add(time.Duration(31-x) * time.Microsecond)
		vec.WithLabelValues("val1", "val2").Set(float64(x))
		x = (x + i) % 23
	}
}
<|endoftext|>


<commit_before>package wrap

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/Songmu/wrapcommander"
	"github.com/mackerelio/mackerel-agent/config"
	"github.com/mackerelio/mkr/logger"
	"golang.org/x/sync/errgroup"
	cli "gopkg.in/urfave/cli.v1"
)

// CommandPlugin is definition of mkr wrap
var Command = cli.Command{
	Name:      "wrap",
	Usage:     "wrap command status",
	ArgsUsage: "[--verbose | -v] [--name | -n <name>] [--memo | -m <memo>] -- /path/to/batch",
	Description: `
    wrap command line
`,
	Action: doWrap,
	Flags: []cli.Flag{
		cli.StringFlag{Name: "name, n", Value: "", Usage: "monitor <name>"},
		cli.BoolFlag{Name: "verbose, v", Usage: "verbose output mode"},
		cli.StringFlag{Name: "memo, m", Value: "", Usage: "monitor <memo>"},
		cli.StringFlag{Name: "H, host", Value: "", Usage: "<hostId>"},
		cli.BoolFlag{Name: "warning, w", Usage: "alert as warning"},
	},
}

func doWrap(c *cli.Context) error {
	confFile := c.GlobalString("conf")
	conf, err := config.LoadConfig(confFile)
	if err != nil {
		return err
	}
	apibase := c.GlobalString("apibase")
	if apibase == "" {
		apibase = conf.Apibase
	}

	apikey := conf.Apikey
	if apikey == "" {
		apikey = os.Getenv("MACKEREL_APIKEY")
	}
	if apikey == "" {
		logger.Log("error", "[mkr wrap] failed to detect Mackerel APIKey. Try to specify in mackerel-agent.conf or export MACKEREL_APIKEY='<Your apikey>'")
	}
	hostID, _ := conf.LoadHostID()
	if c.String("host") != "" {
		hostID = c.String("host")
	}
	if hostID == "" {
		logger.Log("error", "[mkr wrap] failed to load hostID. Try to specify -host option explicitly")
	}
	// Since command execution has the highest priority, even when apikey or
	// hostID is empty, we don't return errors and only output the log here.

	cmd := c.Args()
	if len(cmd) > 0 && cmd[0] == "--" {
		cmd = cmd[1:]
	}
	if len(cmd) < 1 {
		return fmt.Errorf("no commands specified")
	}

	return (&app{
		apibase: apibase,
		name:    c.String("name"),
		verbose: c.Bool("verbose"),
		memo:    c.String("memo"),
		warning: c.Bool("warning"),
		hostID:  hostID,
		apikey:  apikey,
		cmd:     cmd,
	}).run()
}

type app struct {
	apibase string
	name    string
	verbose bool
	memo    string
	warning bool
	hostID  string
	apikey  string
	cmd     []string
}

type result struct {
	Cmd        []string
	Name, Memo string

	Output, Stdout, Stderr string `json:"-"`
	Pid                    int
	ExitCode               *int
	Signaled               bool
	StartAt, EndAt         time.Time

	Msg     string
	Success bool
}

var reg = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizeName(devName string) string {
	return reg.ReplaceAllString(strings.TrimSpace(devName), "_")
}

func (re *result) checkName() string {
	if re.Name != "" {
		return re.Name
	}
	sum := md5.Sum([]byte(strings.Join(re.Cmd, " ")))
	return fmt.Sprintf("mkrwrap-%s-%x",
		normalizeName(filepath.Base(re.Cmd[0])),
		sum[0:3])
}

func (re *result) resultFile() string {
	return filepath.Join(os.TempDir(), fmt.Sprintf("mkrwrap-%s.json", re.checkName()))
}

func (re *result) loadLastResult() (*result, error) {
	prevRe := &result{}
	fname := re.resultFile()

	f, err := os.Open(fname)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	defer f.Close()

	err = json.NewDecoder(f).Decode(prevRe)
	return prevRe, err
}

func (re *result) saveResult() error {
	fname := re.resultFile()
	tmpf, err := ioutil.TempFile(filepath.Dir(fname), "tmp-mkrwrap")
	if err != nil {
		return err
	}
	defer func(tmpfname string) {
		tmpf.Close()
		os.Remove(tmpfname)
	}(tmpf.Name())

	if err := json.NewEncoder(tmpf).Encode(re); err != nil {
		return err
	}
	if err := tmpf.Close(); err != nil {
		return err
	}
	return os.Rename(tmpf.Name(), fname)
}

func (ap *app) run() error {
	re := ap.runCmd()
	return ap.report(re)
}

func (ap *app) runCmd() *result {
	cmd := exec.Command(ap.cmd[0], ap.cmd[1:]...)
	re := &result{
		Cmd:  ap.cmd,
		Name: ap.name,
		Memo: ap.memo,
	}

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		re.Msg = fmt.Sprintf("command invocation failed with follwing error: %s", err)
		return re
	}
	defer stdoutPipe.Close()

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		re.Msg = fmt.Sprintf("command invocation failed with follwing error: %s", err)
		return re
	}
	defer stderrPipe.Close()

	var (
		bufStdout = &bytes.Buffer{}
		bufStderr = &bytes.Buffer{}
		bufMerged = &bytes.Buffer{}
	)
	stdoutPipe2 := io.TeeReader(stdoutPipe, io.MultiWriter(bufStdout, bufMerged))
	stderrPipe2 := io.TeeReader(stderrPipe, io.MultiWriter(bufStderr, bufMerged))

	re.StartAt = time.Now()
	err = cmd.Start()
	if err != nil {
		re.Msg = fmt.Sprintf("command invocation failed with follwing error: %s", err)
		return re
	}
	re.Pid = cmd.Process.Pid
	eg := &errgroup.Group{}

	eg.Go(func() error {
		defer stdoutPipe.Close()
		_, err := io.Copy(os.Stdout, stdoutPipe2)
		return err
	})
	eg.Go(func() error {
		defer stderrPipe.Close()
		_, err := io.Copy(os.Stderr, stderrPipe2)
		return err
	})
	eg.Wait()

	cmdErr := cmd.Wait()
	re.EndAt = time.Now()
	ex := wrapcommander.ResolveExitCode(cmdErr)
	re.ExitCode = &ex
	if *re.ExitCode > 128 {
		w, ok := wrapcommander.ErrorToWaitStatus(cmdErr)
		if ok {
			re.Signaled = w.Signaled()
		}
	}
	if !re.Signaled {
		re.Msg = fmt.Sprintf("command exited with code: %d", *re.ExitCode)
	} else {
		re.Msg = fmt.Sprintf("command died with signal: %d", *re.ExitCode&127)
	}
	re.Stdout = bufStdout.String()
	re.Stderr = bufStderr.String()
	re.Output = bufMerged.String()

	re.Success = *re.ExitCode == 0
	return re
}

func (ap *app) report(re *result) error {
	if ap.apikey == "" || ap.hostID == "" {
		return fmt.Errorf("Both of apikey and hostID are needed to report result to Mackerel")
	}
	lastRe, err := re.loadLastResult()
	if err != nil {
		// resultFile something went wrong.
		// It may be no permission, broken json, not a normal file, and so on.
		// Though it rough, try to delete as workaround
		err := os.RemoveAll(re.resultFile())
		if err != nil {
			// XXX report result here?
			return err
		}
	}
	if lastRe == nil || !re.Success {
		ap.doReport(re)
	}
	return re.saveResult()
}

func (ap *app) doReport(re *result) error {
	return nil
}
<commit_msg>fix reporting condition<commit_after>package wrap

import (
	"bytes"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/Songmu/wrapcommander"
	"github.com/mackerelio/mackerel-agent/config"
	"github.com/mackerelio/mkr/logger"
	"golang.org/x/sync/errgroup"
	cli "gopkg.in/urfave/cli.v1"
)

// CommandPlugin is definition of mkr wrap
var Command = cli.Command{
	Name:      "wrap",
	Usage:     "wrap command status",
	ArgsUsage: "[--verbose | -v] [--name | -n <name>] [--memo | -m <memo>] -- /path/to/batch",
	Description: `
    wrap command line
`,
	Action: doWrap,
	Flags: []cli.Flag{
		cli.StringFlag{Name: "name, n", Value: "", Usage: "monitor <name>"},
		cli.BoolFlag{Name: "verbose, v", Usage: "verbose output mode"},
		cli.StringFlag{Name: "memo, m", Value: "", Usage: "monitor <memo>"},
		cli.StringFlag{Name: "H, host", Value: "", Usage: "<hostId>"},
		cli.BoolFlag{Name: "warning, w", Usage: "alert as warning"},
	},
}

func doWrap(c *cli.Context) error {
	confFile := c.GlobalString("conf")
	conf, err := config.LoadConfig(confFile)
	if err != nil {
		return err
	}
	apibase := c.GlobalString("apibase")
	if apibase == "" {
		apibase = conf.Apibase
	}

	apikey := conf.Apikey
	if apikey == "" {
		apikey = os.Getenv("MACKEREL_APIKEY")
	}
	if apikey == "" {
		logger.Log("error", "[mkr wrap] failed to detect Mackerel APIKey. Try to specify in mackerel-agent.conf or export MACKEREL_APIKEY='<Your apikey>'")
	}
	hostID, _ := conf.LoadHostID()
	if c.String("host") != "" {
		hostID = c.String("host")
	}
	if hostID == "" {
		logger.Log("error", "[mkr wrap] failed to load hostID. Try to specify -host option explicitly")
	}
	// Since command execution has the highest priority, even when apikey or
	// hostID is empty, we don't return errors and only output the log here.

	cmd := c.Args()
	if len(cmd) > 0 && cmd[0] == "--" {
		cmd = cmd[1:]
	}
	if len(cmd) < 1 {
		return fmt.Errorf("no commands specified")
	}

	return (&app{
		apibase: apibase,
		name:    c.String("name"),
		verbose: c.Bool("verbose"),
		memo:    c.String("memo"),
		warning: c.Bool("warning"),
		hostID:  hostID,
		apikey:  apikey,
		cmd:     cmd,
	}).run()
}

type app struct {
	apibase string
	name    string
	verbose bool
	memo    string
	warning bool
	hostID  string
	apikey  string
	cmd     []string
}

type result struct {
	Cmd        []string
	Name, Memo string

	Output, Stdout, Stderr string `json:"-"`
	Pid                    int
	ExitCode               *int
	Signaled               bool
	StartAt, EndAt         time.Time

	Msg     string
	Success bool
}

var reg = regexp.MustCompile(`[^-a-zA-Z0-9_]`)

func normalizeName(devName string) string {
	return reg.ReplaceAllString(strings.TrimSpace(devName), "_")
}

func (re *result) checkName() string {
	if re.Name != "" {
		return re.Name
	}
	sum := md5.Sum([]byte(strings.Join(re.Cmd, " ")))
	return fmt.Sprintf("mkrwrap-%s-%x",
		normalizeName(filepath.Base(re.Cmd[0])),
		sum[0:3])
}

func (re *result) resultFile() string {
	return filepath.Join(os.TempDir(), fmt.Sprintf("mkrwrap-%s.json", re.checkName()))
}

func (re *result) loadLastResult() (*result, error) {
	prevRe := &result{}
	fname := re.resultFile()

	f, err := os.Open(fname)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, nil
		}
		return nil, err
	}
	defer f.Close()

	err = json.NewDecoder(f).Decode(prevRe)
	return prevRe, err
}

func (re *result) saveResult() error {
	fname := re.resultFile()
	tmpf, err := ioutil.TempFile(filepath.Dir(fname), "tmp-mkrwrap")
	if err != nil {
		return err
	}
	defer func(tmpfname string) {
		tmpf.Close()
		os.Remove(tmpfname)
	}(tmpf.Name())

	if err := json.NewEncoder(tmpf).Encode(re); err != nil {
		return err
	}
	if err := tmpf.Close(); err != nil {
		return err
	}
	return os.Rename(tmpf.Name(), fname)
}

func (ap *app) run() error {
	re := ap.runCmd()
	return ap.report(re)
}

func (ap *app) runCmd() *result {
	cmd := exec.Command(ap.cmd[0], ap.cmd[1:]...)
	re := &result{
		Cmd:  ap.cmd,
		Name: ap.name,
		Memo: ap.memo,
	}

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		re.Msg = fmt.Sprintf("command invocation failed with follwing error: %s", err)
		return re
	}
	defer stdoutPipe.Close()

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		re.Msg = fmt.Sprintf("command invocation failed with follwing error: %s", err)
		return re
	}
	defer stderrPipe.Close()

	var (
		bufStdout = &bytes.Buffer{}
		bufStderr = &bytes.Buffer{}
		bufMerged = &bytes.Buffer{}
	)
	stdoutPipe2 := io.TeeReader(stdoutPipe, io.MultiWriter(bufStdout, bufMerged))
	stderrPipe2 := io.TeeReader(stderrPipe, io.MultiWriter(bufStderr, bufMerged))

	re.StartAt = time.Now()
	err = cmd.Start()
	if err != nil {
		re.Msg = fmt.Sprintf("command invocation failed with follwing error: %s", err)
		return re
	}
	re.Pid = cmd.Process.Pid
	eg := &errgroup.Group{}

	eg.Go(func() error {
		defer stdoutPipe.Close()
		_, err := io.Copy(os.Stdout, stdoutPipe2)
		return err
	})
	eg.Go(func() error {
		defer stderrPipe.Close()
		_, err := io.Copy(os.Stderr, stderrPipe2)
		return err
	})
	eg.Wait()

	cmdErr := cmd.Wait()
	re.EndAt = time.Now()
	ex := wrapcommander.ResolveExitCode(cmdErr)
	re.ExitCode = &ex
	if *re.ExitCode > 128 {
		w, ok := wrapcommander.ErrorToWaitStatus(cmdErr)
		if ok {
			re.Signaled = w.Signaled()
		}
	}
	if !re.Signaled {
		re.Msg = fmt.Sprintf("command exited with code: %d", *re.ExitCode)
	} else {
		re.Msg = fmt.Sprintf("command died with signal: %d", *re.ExitCode&127)
	}
	re.Stdout = bufStdout.String()
	re.Stderr = bufStderr.String()
	re.Output = bufMerged.String()

	re.Success = *re.ExitCode == 0
	return re
}

func (ap *app) report(re *result) error {
	if ap.apikey == "" || ap.hostID == "" {
		return fmt.Errorf("Both of apikey and hostID are needed to report result to Mackerel")
	}
	lastRe, err := re.loadLastResult()
	if err != nil {
		// resultFile something went wrong.
		// It may be no permission, broken json, not a normal file, and so on.
		// Though it rough, try to delete as workaround
		err := os.RemoveAll(re.resultFile())
		if err != nil {
			// XXX report result here?
			return err
		}
	}
	if lastRe == nil || !lastRe.Success || !re.Success {
		ap.doReport(re)
	}
	return re.saveResult()
}

func (ap *app) doReport(re *result) error {
	return nil
}
<|endoftext|>


<commit_before>// implementation of ajvondrak's xand core
package xand

import (
	"fmt"
	"github.com/DrItanium/cores/registration/machine"
	"github.com/DrItanium/cores/registration/parser"
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"
)

func RegistrationName() string {
	return "xand"
}

func Register() {}

type MachineRegistrar func(...interface{}) (machine.Machine, error)

func (this MachineRegistrar) New(args ...interface{}) (machine.Machine, error) {
	return this(args)
}

func generateCore(a ...interface{}) (machine.Machine, error) {
	return New()
}

func init() {
	machine.Register(RegistrationName(), MachineRegistrar(generateCore))
}

type Core struct {
	pc     int8
	ir     [3]int8
	memory [128]int8
}

func (this *Core) fetch() bool {
	if (this.pc < 0) || (int(this.pc+2) >= len(this.memory)) {
		return false
	} else {
		this.ir[0] = this.memory[this.pc]
		this.ir[1] = this.memory[this.pc+1]
		this.ir[2] = this.memory[this.pc+2]
		return this.ir[0] >= 0 && this.ir[1] >= 0 && this.ir[2] >= 0
	}
}

func (this *Core) Run() error {
	for this.fetch() {
		// the xand operation it self
		this.memory[this.ir[0]] = this.memory[this.ir[0]] - this.memory[this.ir[1]]
		if this.memory[this.ir[0]] <= 0 {
			this.pc = this.ir[2]
		} else {
			this.pc += 3
		}
	}
	return nil
}

func (this *Core) Startup() error {
	return nil
}

func (this *Core) Shutdown() error {
	return nil
}

func (this *Core) GetDebugStatus() bool {
	return false
}

func (this *Core) SetDebug(_ bool) {

}

func (this *Core) InstallProgram(input <-chan byte) error {
	// read 128 bytes
	for i := 0; i < 128; i++ {
		if value, more := <-input; !more {
			return fmt.Errorf("Not a complete xand memory image")
		} else {
			this.memory[i] = int8(value)
		}
	}
	return nil
}

func (this *Core) Dump(output chan<- byte) error {
	for _, value := range this.memory {
		output <- byte(value)
	}
	return nil
}

func New() (*Core, error) {
	return &Core{}, nil
}

type ParsingRegistrar func(...interface{}) (parser.Parser, error)

func (this ParsingRegistrar) New(args ...interface{}) (parser.Parser, error) {
	return this(args)
}

func generateParser(a ...interface{}) (parser.Parser, error) {
	var p _parser
	if core, err := New(); err != nil {
		return nil, err
	} else {
		p.core = core
		p.labels = make(map[string]int8)
		return &p, nil
	}
}

func init() {
	parser.Register(RegistrationName(), ParsingRegistrar(generateParser))
}

type deferredAddress struct {
	addr  int8
	title string
}

type _parser struct {
	core       *Core
	labels     map[string]int8
	statements []*statement
	deferred   []deferredAddress
}

func (this *_parser) Dump(pipe chan<- byte) error {
	return this.core.Dump(pipe)
}

func (this *_parser) Parse(lines <-chan parser.Entry) error {
	for line := range lines {
		stmt := carveLine(line.Line)
		stmt.index = line.Index
		this.statements = append(this.statements, stmt)
		for _, str := range stmt.contents {
			if err := str.Parse(); err != nil {
				return fmt.Errorf("Error: line: %d : %s\n", line.Index, err)
			}
		}
	}
	return nil
}

type nodeType int

func (this nodeType) String() string {
	switch this {
	case typeId:
		return "id"
	case typeImmediate:
		return "immediate"
	case typeLabel:
		return "label"
	case typeComment:
		return "comment"
	case keywordXand:
		return "xand"
	case keywordDotDotDot:
		return "..."
	default:
		return fmt.Sprintf("%d", this)
	}
}
func (this nodeType) immediate() bool {
	return this == typeImmediate
}
func (this nodeType) comment() bool {
	return this == typeComment
}

const (
	typeId nodeType = iota
	typeImmediate
	typeLabel
	typeComment
	keywordXand
	keywordDotDotDot
)

type node struct {
	Value interface{}
	Type  nodeType
}

func parseDecimalImmediate(str string) (int8, error) {
	val, err := strconv.ParseInt(str, 10, 8)
	return int8(val), err
}

func (this *node) parseLabel(val string) error {
	nVal := strings.TrimSuffix(val, ":")
	q, _ := utf8.DecodeRuneInString(nVal)
	if !unicode.IsLetter(q) {
		return fmt.Errorf("Label %s starts with a non letter %s!", nVal, q)
	} else {
		this.Type = typeLabel
		this.Value = nVal
		// now parse the label as a entirely new node and see if we get a register back
		if nVal == "xand" {
			return fmt.Errorf("Can't name a label xand")
		} else {
			return nil
		}
	}
}

func (this *node) parseImmediate(val string) error {
	this.Type = typeImmediate
	if v, err := parseDecimalImmediate(val[1:]); err != nil {
		return err
	} else {
		this.Value = v
		return nil
	}
}

var keywords = map[string]nodeType{
	"xand": keywordXand,
	"...":  keywordDotDotDot,
}

func (this *node) parseGeneric(val string) error {
	if v, ok := keywords[val]; ok {
		this.Type = v
		return nil
	} else {
		return fmt.Errorf("Unknown statement %s", val)
	}
}

func (this *node) Parse() error {
	if this.Type == typeId {
		val := this.Value.(string)
		if this.parseGeneric(val) == nil {

		} else if strings.HasSuffix(val, ":") {
			return this.parseLabel(val)
		} else if strings.HasPrefix(val, ";") {
			this.Type = typeComment
			this.Value = strings.TrimPrefix(val, ";")
		} else if strings.HasPrefix(val, "#") {
			return this.parseImmediate(val)
		}
	}
	return nil
}

func (this *node) isComment() bool {
	return this.Type == typeComment
}

func (this *node) isLabel() bool {
	return this.Type == typeLabel
}

type statement struct {
	contents []*node
	index    int
}

func (this *statement) Add(value string, t nodeType) {
	// always trim before adding
	str := strings.TrimSpace(value)
	if len(str) > 0 {
		this.contents = append(this.contents, &node{Value: str, Type: t})
	}
}
func (this *statement) AddUnknown(value string) {
	this.Add(value, typeId)
}
func (this *statement) String() string {
	str := fmt.Sprintf("%d: ", this.index)
	for _, n := range this.contents {
		str += fmt.Sprintf(" %T: %s ", n, *n)
	}
	return str
}
func (this *statement) First() (*node, error) {
	if len(this.contents) == 0 {
		return nil, fmt.Errorf("Empty statement!")
	} else {
		return this.contents[0], nil
	}
}
func (this *statement) Rest() []*node {
	return this.contents[1:]
}

func carveLine(line string) *statement {
	// trim the damn line first
	data := strings.TrimSpace(line)
	var s statement
	if len(data) == 0 {
		return &s
	}
	oldStart := 0
	start := 0
	// skip the strings at the beginning
	for width := 0; start < len(data); start += width {
		var r rune
		next := data[start:]
		r, width = utf8.DecodeRuneInString(next)
		if unicode.IsSpace(r) {
			s.AddUnknown(data[oldStart:start])
			oldStart = start
		} else if r == ';' {
			// consume the rest of the data
			s.AddUnknown(data[oldStart:start])
			// then capture the comment
			s.Add(data[start:], typeComment)
			oldStart = start
			break
		}
	}
	if oldStart < start {
		s.AddUnknown(data[oldStart:])
	}
	return &s
}

func (this *_parser) Process() error {
	for _, stmt := range this.statements {
		if err := this.parseStatement(stmt); err != nil {
			return fmt.Errorf("Error: line %d: msg: %s", stmt.index, err)
		}
	}
	for _, d := range this.deferred {
		if entry, ok := this.labels[d.title]; !ok {
			return fmt.Errorf("Label %s not defined!", d.title)
		} else {
			this.core.memory[d.addr] = entry
		}
	}
	return nil
}

func (this *_parser) newLabel(n *node) error {
	name := n.Value.(string)
	if _, ok := this.labels[name]; ok {
		return fmt.Errorf("Label %s is already defined!", name)
	} else {
		this.labels[name] = this.core.pc
		return nil
	}
}

func (this *_parser) parseStatement(stmt *statement) error {
	first, err := stmt.First()
	if err != nil {
		return err
	}
	rest := stmt.Rest()
	switch first.Type {
	case typeComment:
		if len(rest) > 0 {
			panic("Programmer Failure! Found something following a comment node in a statement. This is impossible!!!!")
		} else {
			return nil
		}
	case typeLabel:
		if err := this.newLabel(first); err != nil {
			return err
		} else if len(rest) > 0 {
			if this.core.pc < 0 {
				return fmt.Errorf("Too many instructions defined!")
			}
			// if there are more entries on the line then check them out
			var s statement
			s.index = stmt.index
			s.contents = rest
			return this.parseStatement(&s)
		}
	case keywordXand:
		if len(rest) == 3 {
			if this.core.pc < 0 {
				return fmt.Errorf("Too many instructions defined!")
			}
			var s statement
			s.index = stmt.index
			s.contents = rest
			return this.parseStatement(&s)
		} else {
			return fmt.Errorf("xand requires three arguments")
		}
	case keywordDotDotDot:
		this.core.memory[this.core.pc] = this.core.pc + 1
		this.core.pc++
		// hmmm should we allow this to continue on?...nope
		if len(rest) > 0 {
			return fmt.Errorf("... has to terminate a statement")
		}
	case typeImmediate:
		// just install the value to the current address
		this.core.memory[this.core.pc] = first.Value.(int8)
		this.core.pc++
		if len(rest) > 0 {
			if this.core.pc < 0 {
				return fmt.Errorf("Too many instructions defined!")
			}
			var s statement
			s.index = stmt.index
			s.contents = rest
			return this.parseStatement(&s)
		}
	case typeId:
		// defer statement for the time being
		if addr, ok := this.labels[first.Value.(string)]; !ok {
			this.deferred = append(this.deferred, deferredAddress{addr: this.core.pc, title: first.Value.(string)})
		} else {
			this.core.memory[this.core.pc] = addr
		}
		this.core.pc++
		if len(rest) > 0 {
			if this.core.pc < 0 {
				return fmt.Errorf("Too many instructions defined!")
			}
			var s statement
			s.index = stmt.index
			s.contents = rest
			return this.parseStatement(&s)
		}
	default:
		return fmt.Errorf("Unhandled nodeType %d: %s", first.Type, first.Value)
	}
	return nil
}
<commit_msg>Added typedefs to xand<commit_after>// implementation of ajvondrak's xand core
package xand

import (
	"fmt"
	"github.com/DrItanium/cores/registration/machine"
	"github.com/DrItanium/cores/registration/parser"
	"strconv"
	"strings"
	"unicode"
	"unicode/utf8"
)

type Word int8

const MemorySize = 128

func RegistrationName() string {
	return "xand"
}

func Register() {}

type MachineRegistrar func(...interface{}) (machine.Machine, error)

func (this MachineRegistrar) New(args ...interface{}) (machine.Machine, error) {
	return this(args)
}

func generateCore(a ...interface{}) (machine.Machine, error) {
	return New()
}

func init() {
	machine.Register(RegistrationName(), MachineRegistrar(generateCore))
}

type Core struct {
	pc     Word
	ir     [3]Word
	memory [MemorySize]Word
}

func (this *Core) fetch() bool {
	if (this.pc < 0) || (int(this.pc+2) >= len(this.memory)) {
		return false
	} else {
		this.ir[0] = this.memory[this.pc]
		this.ir[1] = this.memory[this.pc+1]
		this.ir[2] = this.memory[this.pc+2]
		return this.ir[0] >= 0 && this.ir[1] >= 0 && this.ir[2] >= 0
	}
}

func (this *Core) Run() error {
	for this.fetch() {
		// the xand operation it self
		this.memory[this.ir[0]] = this.memory[this.ir[0]] - this.memory[this.ir[1]]
		if this.memory[this.ir[0]] <= 0 {
			this.pc = this.ir[2]
		} else {
			this.pc += 3
		}
	}
	return nil
}

func (this *Core) Startup() error {
	return nil
}

func (this *Core) Shutdown() error {
	return nil
}

func (this *Core) GetDebugStatus() bool {
	return false
}

func (this *Core) SetDebug(_ bool) {

}

func (this *Core) InstallProgram(input <-chan byte) error {
	// read 128 bytes
	for i := 0; i < MemorySize; i++ {
		if value, more := <-input; !more {
			return fmt.Errorf("Not a complete xand memory image")
		} else {
			this.memory[i] = Word(value)
		}
	}
	return nil
}

func (this *Core) Dump(output chan<- byte) error {
	for _, value := range this.memory {
		output <- byte(value)
	}
	return nil
}

func New() (*Core, error) {
	return &Core{}, nil
}

type ParsingRegistrar func(...interface{}) (parser.Parser, error)

func (this ParsingRegistrar) New(args ...interface{}) (parser.Parser, error) {
	return this(args)
}

func generateParser(a ...interface{}) (parser.Parser, error) {
	var p _parser
	if core, err := New(); err != nil {
		return nil, err
	} else {
		p.core = core
		p.labels = make(map[string]Word)
		return &p, nil
	}
}

func init() {
	parser.Register(RegistrationName(), ParsingRegistrar(generateParser))
}

type deferredAddress struct {
	addr  Word
	title string
}

type _parser struct {
	core       *Core
	labels     map[string]Word
	statements []*statement
	deferred   []deferredAddress
}

func (this *_parser) Dump(pipe chan<- byte) error {
	return this.core.Dump(pipe)
}

func (this *_parser) Parse(lines <-chan parser.Entry) error {
	for line := range lines {
		stmt := carveLine(line.Line)
		stmt.index = line.Index
		this.statements = append(this.statements, stmt)
		for _, str := range stmt.contents {
			if err := str.Parse(); err != nil {
				return fmt.Errorf("Error: line: %d : %s\n", line.Index, err)
			}
		}
	}
	return nil
}

type nodeType int

func (this nodeType) String() string {
	switch this {
	case typeId:
		return "id"
	case typeImmediate:
		return "immediate"
	case typeLabel:
		return "label"
	case typeComment:
		return "comment"
	case keywordXand:
		return "xand"
	case keywordDotDotDot:
		return "..."
	default:
		return fmt.Sprintf("%d", this)
	}
}
func (this nodeType) immediate() bool {
	return this == typeImmediate
}
func (this nodeType) comment() bool {
	return this == typeComment
}

const (
	typeId nodeType = iota
	typeImmediate
	typeLabel
	typeComment
	keywordXand
	keywordDotDotDot
)

type node struct {
	Value interface{}
	Type  nodeType
}

func parseDecimalImmediate(str string) (Word, error) {
	val, err := strconv.ParseInt(str, 10, 8)
	return Word(val), err
}

func (this *node) parseLabel(val string) error {
	nVal := strings.TrimSuffix(val, ":")
	q, _ := utf8.DecodeRuneInString(nVal)
	if !unicode.IsLetter(q) {
		return fmt.Errorf("Label %s starts with a non letter %s!", nVal, q)
	} else {
		this.Type = typeLabel
		this.Value = nVal
		// now parse the label as a entirely new node and see if we get a register back
		if nVal == "xand" {
			return fmt.Errorf("Can't name a label xand")
		} else {
			return nil
		}
	}
}

func (this *node) parseImmediate(val string) error {
	this.Type = typeImmediate
	if v, err := parseDecimalImmediate(val[1:]); err != nil {
		return err
	} else {
		this.Value = v
		return nil
	}
}

var keywords = map[string]nodeType{
	"xand": keywordXand,
	"...":  keywordDotDotDot,
}

func (this *node) parseGeneric(val string) error {
	if v, ok := keywords[val]; ok {
		this.Type = v
		return nil
	} else {
		return fmt.Errorf("Unknown statement %s", val)
	}
}

func (this *node) Parse() error {
	if this.Type == typeId {
		val := this.Value.(string)
		if this.parseGeneric(val) == nil {

		} else if strings.HasSuffix(val, ":") {
			return this.parseLabel(val)
		} else if strings.HasPrefix(val, ";") {
			this.Type = typeComment
			this.Value = strings.TrimPrefix(val, ";")
		} else if strings.HasPrefix(val, "#") {
			return this.parseImmediate(val)
		}
	}
	return nil
}

func (this *node) isComment() bool {
	return this.Type == typeComment
}

func (this *node) isLabel() bool {
	return this.Type == typeLabel
}

type statement struct {
	contents []*node
	index    int
}

func (this *statement) Add(value string, t nodeType) {
	// always trim before adding
	str := strings.TrimSpace(value)
	if len(str) > 0 {
		this.contents = append(this.contents, &node{Value: str, Type: t})
	}
}
func (this *statement) AddUnknown(value string) {
	this.Add(value, typeId)
}
func (this *statement) String() string {
	str := fmt.Sprintf("%d: ", this.index)
	for _, n := range this.contents {
		str += fmt.Sprintf(" %T: %s ", n, *n)
	}
	return str
}
func (this *statement) First() (*node, error) {
	if len(this.contents) == 0 {
		return nil, fmt.Errorf("Empty statement!")
	} else {
		return this.contents[0], nil
	}
}
func (this *statement) Rest() []*node {
	return this.contents[1:]
}

func carveLine(line string) *statement {
	// trim the damn line first
	data := strings.TrimSpace(line)
	var s statement
	if len(data) == 0 {
		return &s
	}
	oldStart := 0
	start := 0
	// skip the strings at the beginning
	for width := 0; start < len(data); start += width {
		var r rune
		next := data[start:]
		r, width = utf8.DecodeRuneInString(next)
		if unicode.IsSpace(r) {
			s.AddUnknown(data[oldStart:start])
			oldStart = start
		} else if r == ';' {
			// consume the rest of the data
			s.AddUnknown(data[oldStart:start])
			// then capture the comment
			s.Add(data[start:], typeComment)
			oldStart = start
			break
		}
	}
	if oldStart < start {
		s.AddUnknown(data[oldStart:])
	}
	return &s
}

func (this *_parser) Process() error {
	for _, stmt := range this.statements {
		if err := this.parseStatement(stmt); err != nil {
			return fmt.Errorf("Error: line %d: msg: %s", stmt.index, err)
		}
	}
	for _, d := range this.deferred {
		if entry, ok := this.labels[d.title]; !ok {
			return fmt.Errorf("Label %s not defined!", d.title)
		} else {
			this.core.memory[d.addr] = entry
		}
	}
	return nil
}

func (this *_parser) newLabel(n *node) error {
	name := n.Value.(string)
	if _, ok := this.labels[name]; ok {
		return fmt.Errorf("Label %s is already defined!", name)
	} else {
		this.labels[name] = this.core.pc
		return nil
	}
}

func (this *_parser) parseStatement(stmt *statement) error {
	first, err := stmt.First()
	if err != nil {
		return err
	}
	rest := stmt.Rest()
	switch first.Type {
	case typeComment:
		if len(rest) > 0 {
			panic("Programmer Failure! Found something following a comment node in a statement. This is impossible!!!!")
		} else {
			return nil
		}
	case typeLabel:
		if err := this.newLabel(first); err != nil {
			return err
		} else if len(rest) > 0 {
			if this.core.pc < 0 {
				return fmt.Errorf("Too many instructions defined!")
			}
			// if there are more entries on the line then check them out
			var s statement
			s.index = stmt.index
			s.contents = rest
			return this.parseStatement(&s)
		}
	case keywordXand:
		if len(rest) == 3 {
			if this.core.pc < 0 {
				return fmt.Errorf("Too many instructions defined!")
			}
			var s statement
			s.index = stmt.index
			s.contents = rest
			return this.parseStatement(&s)
		} else {
			return fmt.Errorf("xand requires three arguments")
		}
	case keywordDotDotDot:
		this.core.memory[this.core.pc] = this.core.pc + 1
		this.core.pc++
		// hmmm should we allow this to continue on?...nope
		if len(rest) > 0 {
			return fmt.Errorf("... has to terminate a statement")
		}
	case typeImmediate:
		// just install the value to the current address
		this.core.memory[this.core.pc] = first.Value.(Word)
		this.core.pc++
		if len(rest) > 0 {
			if this.core.pc < 0 {
				return fmt.Errorf("Too many instructions defined!")
			}
			var s statement
			s.index = stmt.index
			s.contents = rest
			return this.parseStatement(&s)
		}
	case typeId:
		// defer statement for the time being
		if addr, ok := this.labels[first.Value.(string)]; !ok {
			this.deferred = append(this.deferred, deferredAddress{addr: this.core.pc, title: first.Value.(string)})
		} else {
			this.core.memory[this.core.pc] = addr
		}
		this.core.pc++
		if len(rest) > 0 {
			if this.core.pc < 0 {
				return fmt.Errorf("Too many instructions defined!")
			}
			var s statement
			s.index = stmt.index
			s.contents = rest
			return this.parseStatement(&s)
		}
	default:
		return fmt.Errorf("Unhandled nodeType %d: %s", first.Type, first.Value)
	}
	return nil
}
<|endoftext|>


<commit_before>package stack

import "sync"

type Stack struct {
	stack []interface{}
	len   int
	lock  sync.Mutex
}

func Constructor() *Stack {
	s := &Stack{}
	s.stack = make([]interface{}, 0)
	s.len = 0

	return s
}

func (s *Stack) Len() int {
	s.lock.Lock()
	defer s.lock.Unlock()

	return s.len
}

func (s *Stack) IsEmpty() bool {
	s.lock.Lock()
	defer s.lock.Unlock()

	return s.len == 0
}

func (s *Stack) Pop() (element interface{}) {
	s.lock.Lock()
	defer s.lock.Unlock()

	element, s.stack = s.stack[0], s.stack[1:]
	s.len--
	return
}

func (s *Stack) Push(element interface{}) {
	s.lock.Lock()
	defer s.lock.Unlock()

	prepend := []interface{}{element}
	s.stack = append(prepend, s.stack...)
	s.len++
}

func (s *Stack) Top() interface{} {
	s.lock.Lock()
	defer s.lock.Unlock()

	return s.stack[0]
}
<commit_msg>use RWMutex instead of Mutex<commit_after>package stack

import "sync"

type Stack struct {
	stack []interface{}
	len   int
	sync.RWMutex
}

func Constructor() *Stack {
	s := &Stack{}
	s.stack = make([]interface{}, 0)
	s.len = 0

	return s
}

func (s *Stack) Len() int {
	s.RLock()
	defer s.RUnlock()

	return s.len
}

func (s *Stack) IsEmpty() bool {
	s.RLock()
	defer s.RUnlock()

	return s.len == 0
}

func (s *Stack) Pop() (element interface{}) {
	s.Lock()
	defer s.Unlock()

	element, s.stack = s.stack[0], s.stack[1:]
	s.len--
	return
}

func (s *Stack) Push(element interface{}) {
	s.Lock()
	defer s.Unlock()

	prepend := []interface{}{element}
	s.stack = append(prepend, s.stack...)
	s.len++
}

func (s *Stack) Top() interface{} {
	s.RLock()
	defer s.RUnlock()

	return s.stack[0]
}
<|endoftext|>


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

package server

import (
	"database/sql"
	_ "github.com/go-sql-driver/mysql"
)

/*
https://code.google.com/p/go-wiki/wiki/SQLDrivers
https://github.com/go-sql-driver/mysql
http://go-database-sql.org/
//
create database pubsubsql;
create user pubsubsql identified by 'pubsubsql';
GRANT ALL PRIVILEGES ON *.* TO 'pubsubsql'@'%' WITH GRANT OPTION;
create table table_a(id int);
create table table_b(id int);
create table table_c(id int);
show tables
 */
type mysqlConnection struct {
	dbConn *sql.DB
	address string
	lastError error
}

func newMysqlConnection() *mysqlConnection {
	return &mysqlConnection {
		dbConn: nil,
		address: "",
		lastError: nil,
	}
}

func (this *mysqlConnection) hasError() bool {
	return nil != this.lastError
}

func (this *mysqlConnection) hasNoError() bool {
	return ! this.hasError()
}

func (this *mysqlConnection) getLastError() string {
	if (this.hasError()) {
		return this.lastError.Error()
	} else {
		return ""
	}
}

func (this *mysqlConnection) isConnected() bool {
	if nil == this.dbConn {
		return false
	} else {
		this.lastError = this.dbConn.Ping()
		if nil != this.lastError {
			this.dbConn.Close()
			this.dbConn = nil
			return false
		}
		return true
	}
}

func (this *mysqlConnection) isDisconnected() bool {
	return ! this.isConnected()
}

func (this *mysqlConnection) disconnect() {
	this.lastError = nil
	if this.isConnected() {
		this.lastError = this.dbConn.Close()
		this.dbConn = nil
	}
}

func (this *mysqlConnection) connect(address string) bool {
	this.lastError = nil
	if this.isDisconnected() {
		this.address = address
		// "pubsubsql:pubsubsql@/pubsubsql"
		this.dbConn, this.lastError = sql.Open("mysql", this.address)
		if nil != this.lastError {
			this.dbConn = nil
			return false
		}
	}
	return this.isConnected();
}

func (this *mysqlConnection) findTables() []string {
	this.lastError = nil
	tables := make([]string, 0)
	if (this.isDisconnected()) {
		return tables
	}
	rows, err := this.dbConn.Query("show tables")
	this.lastError = err
	if nil != this.lastError {
		logError(this.lastError)
		return tables
	}
	tableName := ""
	for rows.Next() {
		this.lastError = rows.Scan(&tableName)
		if  nil != this.lastError {
			logError(this.lastError)
			return tables
		}
		tables = append(tables, tableName)
	}
	this.lastError = rows.Err()
	if nil != this.lastError {
		logError(this.lastError)
		return tables
	}
	//
	return tables
}
<commit_msg>return "not connected" from mysql find tables<commit_after>/* Copyright (C) 2013 CompleteDB LLC.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with PubSubSQL.  If not, see <http://www.gnu.org/licenses/>.
 */

package server

import (
	"database/sql"
	_ "github.com/go-sql-driver/mysql"
)

/*
https://code.google.com/p/go-wiki/wiki/SQLDrivers
https://github.com/go-sql-driver/mysql
http://go-database-sql.org/
//
create database pubsubsql;
create user pubsubsql identified by 'pubsubsql';
GRANT ALL PRIVILEGES ON *.* TO 'pubsubsql'@'%' WITH GRANT OPTION;
create table table_a(id int);
create table table_b(id int);
create table table_c(id int);
show tables
 */
type mysqlConnection struct {
	dbConn *sql.DB
	address string
	lastError string
}

func newMysqlConnection() *mysqlConnection {
	return &mysqlConnection {
		dbConn: nil,
		address: "",
		lastError: "",
	}
}

func (this *mysqlConnection) hasError() bool {
	return "" != this.lastError
}

func (this *mysqlConnection) hasNoError() bool {
	return ! this.hasError()
}

func (this *mysqlConnection) getLastError() string {
	return this.lastError
}

func (this *mysqlConnection) isConnected() bool {
	if nil == this.dbConn {
		return false
	} else {
		err := this.dbConn.Ping()
		if nil != err {
			this.lastError = err.Error()
			this.dbConn.Close()
			this.dbConn = nil
			logError(this.lastError)
			return false
		}
		return true
	}
}

func (this *mysqlConnection) isDisconnected() bool {
	return ! this.isConnected()
}

func (this *mysqlConnection) disconnect() {
	this.lastError = ""
	if this.isConnected() {
		err := this.dbConn.Close()
		if nil != err {
			this.lastError = err.Error()
			logError(this.lastError)
		}
		this.dbConn = nil
	}
}

func (this *mysqlConnection) connect(address string) bool {
	this.lastError = ""
	if this.isDisconnected() {
		this.address = address
		// "pubsubsql:pubsubsql@/pubsubsql"
		var err error
		this.dbConn, err = sql.Open("mysql", this.address)
		if nil != err {
			this.lastError = err.Error()
			this.dbConn = nil
			logError(this.lastError)
			return false
		}
	}
	return this.isConnected();
}

func (this *mysqlConnection) findTables() []string {
	this.lastError = ""
	tables := make([]string, 0)
	if (this.isDisconnected()) {
		this.lastError = "not connected to mysql"
		return tables
	}
	rows, err := this.dbConn.Query("show tables")
	if nil != err {
		this.lastError = err.Error()
		logError(this.lastError)
		return tables
	}
	tableName := ""
	for rows.Next() {
		err := rows.Scan(&tableName)
		if  nil != err {
			this.lastError = err.Error()
			logError(this.lastError)
			return tables
		}
		tables = append(tables, tableName)
	}
	err = rows.Err()
	if nil != err {
		this.lastError = err.Error()
		logError(this.lastError)
		return tables
	}
	//
	return tables
}
<|endoftext|>


<commit_before>package zssh

import (
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/yuin/gopher-lua"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// system configurations.
var (
	SystemWideConfigFile string
	PerUserConfigFile    string
	Version              = "0.5.0"
)

var IgnoreError flag.ErrorHandling = 9999

func Start() error {
	var printFlag, configFlag, systemConfigFlag, debugFlag, hostsFlag, verboseFlag, tagsFlag, zshCompletinFlag bool
	var configFile, shellPath string
	var rsyncArg string
	filters := []string{}

	if len(os.Args) == 1 {
		printUsage()
		return nil
	}

	args := os.Args[1:]
	for {
		if len(args) == 0 {
			break
		}

		arg := args[0]
		if arg == "--print" {
			printFlag = true
		} else if arg == "--config" {
			configFlag = true
		} else if arg == "--system-config" {
			systemConfigFlag = true
		} else if arg == "--debug" {
			debugFlag = true
		} else if arg == "--hosts" {
			hostsFlag = true
		} else if arg == "--verbose" {
			verboseFlag = true
		} else if arg == "--filter" {
			if len(args) < 2 {
				return fmt.Errorf("--filter reguires an argument.")
			}
			filters = append(filters, args[1])
			args = args[1:]
		} else if strings.HasPrefix(arg, "--filter=") {
			filters = append(filters, strings.Split(arg, "=")[1])
		} else if arg == "--tags" {
			tagsFlag = true
		} else if arg == "--zsh-completion" {
			zshCompletinFlag = true
		} else if arg == "--config-file" {
			if len(args) < 2 {
				return fmt.Errorf("--config-file reguires an argument.")
			}
			configFile = args[1]
			args = args[1:]
		} else if strings.HasPrefix(arg, "--config-file=") {
			configFile = strings.Split(arg, "=")[1]
		} else if arg == "--shell" {
			if len(args) < 2 {
				return fmt.Errorf("--shell reguires an argument.")
			}
			shellPath = args[1]
			args = args[1:]
		} else if strings.HasPrefix(arg, "--shell=") {
			shellPath = strings.Split(arg, "=")[1]
		} else if arg == "--rsync" {
			if len(args) < 2 {
				return fmt.Errorf("--rsync reguires an argument.")
			}
			rsyncArg = args[1]
			args = args[1:]
		} else if strings.HasPrefix(arg, "--rsync=") {
			rsyncArg = strings.Split(arg, "=")[1]
		} else {
			break
		}

		args = args[1:]
	}

	if zshCompletinFlag {
		fmt.Print(ZSH_COMPLETION)
		return nil
	}

	if configFlag {
		ShellExec("$EDITOR " + PerUserConfigFile)
		return nil
	}

	if systemConfigFlag {
		ShellExec("$EDITOR " + SystemWideConfigFile)
		return nil
	}

	var shellContent []byte

	if shellPath != "" {
		if strings.HasPrefix(shellPath, "http://") || strings.HasPrefix(shellPath, "https://") {
			// get script from remote using http.
			if debugFlag {
				fmt.Printf("[zssh debug] get script using http from '%s'\n", shellPath)
			}

			var httpClient *http.Client
			if strings.HasPrefix(shellPath, "https://") {
				tr := &http.Transport{
					TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
				}
				httpClient = &http.Client{Transport: tr}
			} else {
				httpClient = &http.Client{}
			}

			resp, err := httpClient.Get(shellPath)
			if err != nil {
				return err
			}
			defer resp.Body.Close()
			b, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return err
			}

			shellContent = b
		} else {
			// get script from the file system.
			b, err := ioutil.ReadFile(shellPath)
			if err != nil {
				return err
			}
			shellContent = b
		}
	}

	// set up the lua state.
	L := lua.NewState()
	defer L.Close()

	// load lua custom functions
	LoadFunctions(L)

	if debugFlag {
		fmt.Printf("[zssh debug] loaded lua functions\n")
	}

	// load specific config file
	if configFile != "" {
		_, err := os.Stat(configFile)
		if err != nil {
			return err
		}

		if err := L.DoFile(configFile); err != nil {
			return err
		}

		if debugFlag {
			fmt.Printf("[zssh debug] loaded config file: %s \n", configFile)
		}

	} else {
		// load system wide config
		if _, err := os.Stat(SystemWideConfigFile); err == nil {
			if err := L.DoFile(SystemWideConfigFile); err != nil {
				return err
			}

			if debugFlag {
				fmt.Printf("[zssh debug] loaded config file: %s \n", SystemWideConfigFile)
			}

		}

		// load per-user wide config
		if _, err := os.Stat(PerUserConfigFile); err == nil {
			if err := L.DoFile(PerUserConfigFile); err != nil {
				return err
			}

			if debugFlag {
				fmt.Printf("[zssh debug] loaded config file: %s \n", PerUserConfigFile)
			}

		}
	}

	// generate ssh hosts config
	content, err := GenHostsConfig()
	if err != nil {
		return err
	}

	// only print generated config
	if printFlag {
		fmt.Println(string(content))
		return nil
	}

	// only print hosts list
	if hostsFlag {
		var hosts []*Host
		if len(filters) > 0 {
			hosts = HostsByTags(filters)
		} else {
			hosts = Hosts
		}

		for _, host := range hosts {
			if !host.Hidden {
				if verboseFlag {
					fmt.Printf("%s\t%s\n", host.Name, host.Description)
				} else {
					fmt.Printf("%s\n", host.Name)
				}
			}
		}

		return nil
	}

	if tagsFlag {
		for _, tag := range Tags() {
			fmt.Printf("%s\n", tag)
		}

		return nil
	}

	// generate temporary ssh config file
	tmpFile, err := ioutil.TempFile("", "zssh.ssh_config.")
	if err != nil {
		return err
	}
	defer func() {
		tmpFile.Close()
		os.Remove(tmpFile.Name())

		if debugFlag {
			fmt.Printf("[zssh debug] deleted config file: %s \n", tmpFile.Name())
		}

	}()
	generatedSSHConfigFile := tmpFile.Name()

	if debugFlag {
		fmt.Printf("[zssh debug] generated config file: %s \n", generatedSSHConfigFile)
	}

	// update temporary sss config file
	err = ioutil.WriteFile(generatedSSHConfigFile, content, 0644)
	if err != nil {
		return err
	}

	if shellPath == "" && rsyncArg == "" {
		// get hooks
		var hooks map[string]func() error

		// Limitation!: hooks fires only when the hostname is specified by the first argument.
		if len(args) > 0 {
			hostname := args[0]
			if host := GetHost(hostname); host != nil {
				hooks = host.Hooks
			}
		}

		// run before hook
		if before := hooks["before"]; before != nil {
			if debugFlag {
				fmt.Printf("[zssh debug] run before hook\n")
			}
			err := before()
			if err != nil {
				return err
			}
		}

		// register after hook
		defer func() {
			// after hook
			if after := hooks["after"]; after != nil {
				if debugFlag {
					fmt.Printf("[zssh debug] run after hook\n")
				}
				err := after()
				if err != nil {
					panic(err)
				}
			}
		}()
	}

	// setup ssh command args
	sshComandArgs := []string{"-F", generatedSSHConfigFile}
	sshComandArgs = append(sshComandArgs, args[:]...)
	if shellPath != "" {
		sshComandArgs = append(sshComandArgs, "bash", "-se")
	}

	if rsyncArg != "" {
		// rsync mode.
		if debugFlag {
			fmt.Printf("[zssh debug] use rsync mode.\n")
		}

		rsyncSSHOption := `-e "ssh ` + strings.Join(sshComandArgs, " ") + `"`
		rsyncCommand := "rsync " + rsyncSSHOption + " " + rsyncArg

		if debugFlag {
			fmt.Printf("[zssh debug] real rsync command: %v\n", rsyncCommand)
		}

		err := ShellExec(rsyncCommand)
		if err != nil {
			return err
		}
	} else {
		// normal ssh mode.

		// execute ssh commmand
		cmd := exec.Command("ssh", sshComandArgs[:]...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr

		if shellPath != "" {
			// input shell script from stdin.
			cmd.Stdin = bytes.NewBuffer(shellContent)

			if debugFlag {
				fmt.Printf("[zssh debug] shell content:\n%s\n", string(shellContent))
			}

		} else {
			cmd.Stdin = os.Stdin
		}

		if debugFlag {
			fmt.Printf("[zssh debug] real ssh command: %v \n", cmd.Args)
		}

		err = cmd.Run()
		if err != nil {
			return err
		}
	}

	return nil
}

func printUsage() {
	// print usage.
	fmt.Println(`Usage: zssh [<options>] [<ssh options and args...>]

zssh is an extended ssh command.
version ` + Version + `

zssh options:
  --print                 Print generated ssh config.
  --config                Edit per-user config file.
  --system-config         Edit system wide config file.
  --config-file <FILE>    Load configuration from the specific file.
  --hosts                 List hosts. This option can use with additional options.
  --filter <TAG>          (Using with --hosts option) Show only the hosts configured with a tag.
  --verbose               (Using with --hosts option) List hosts with description.
  --tags                  List tags.
  --zsh-completion        Output zsh completion code.
  --debug                 Output debug log

zssh options for convenient functionalities:
  --shell <PATH> <HOSTNAME>              Execute a shell script of the path on the remote host.
  --rsync <rsync options and args...>    Execute rsync using zssh config.

And the following is original ssh command usage...
`)
	// show ssh help
	cmd := exec.Command("ssh")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	cmd.Run()
}

func init() {
	if SystemWideConfigFile == "" {
		SystemWideConfigFile = "/etc/zssh/config.lua"
	}
	if PerUserConfigFile == "" {
		home := userHomeDir()
		PerUserConfigFile = filepath.Join(home, ".zssh/config.lua")
	}
}

var ZSH_COMPLETION = `
_zssh_hosts() {
    local -a __zssh_hosts
    PRE_IFS=$IFS
    IFS=$'\n'
    __zssh_hosts=($(zssh --hosts --verbose | awk -F'\t' '{print $1":"$2}'))
    IFS=$PRE_IFS
    _describe -t host "host" __zssh_hosts
}

_zssh () {
    local curcontext="$curcontext" state line
    typeset -A opt_args

    _arguments \
        '1: :->command'

    case $state in
        command)
            _zssh_hosts
            ;;
        *)
            _files
            ;;
    esac
}

compdef _zssh zssh

`
<commit_msg>tiny update<commit_after>package zssh

import (
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/yuin/gopher-lua"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

// system configurations.
var (
	SystemWideConfigFile string
	PerUserConfigFile    string
	Version              = "0.5.0"
)

var IgnoreError flag.ErrorHandling = 9999

func Start() error {
	var printFlag, configFlag, systemConfigFlag, debugFlag, hostsFlag, verboseFlag, tagsFlag, zshCompletinFlag, bashCompletinFlag bool
	var configFile, shellPath string
	var rsyncArg string
	filters := []string{}

	if len(os.Args) == 1 {
		printUsage()
		return nil
	}

	args := os.Args[1:]
	for {
		if len(args) == 0 {
			break
		}

		arg := args[0]
		if arg == "--print" {
			printFlag = true
		} else if arg == "--config" {
			configFlag = true
		} else if arg == "--system-config" {
			systemConfigFlag = true
		} else if arg == "--debug" {
			debugFlag = true
		} else if arg == "--hosts" {
			hostsFlag = true
		} else if arg == "--verbose" {
			verboseFlag = true
		} else if arg == "--filter" {
			if len(args) < 2 {
				return fmt.Errorf("--filter reguires an argument.")
			}
			filters = append(filters, args[1])
			args = args[1:]
		} else if strings.HasPrefix(arg, "--filter=") {
			filters = append(filters, strings.Split(arg, "=")[1])
		} else if arg == "--tags" {
			tagsFlag = true
		} else if arg == "--zsh-completion" {
			zshCompletinFlag = true
		} else if arg == "--bash-completion" {
			bashCompletinFlag = true
		} else if arg == "--config-file" {
			if len(args) < 2 {
				return fmt.Errorf("--config-file reguires an argument.")
			}
			configFile = args[1]
			args = args[1:]
		} else if strings.HasPrefix(arg, "--config-file=") {
			configFile = strings.Split(arg, "=")[1]
		} else if arg == "--shell" {
			if len(args) < 2 {
				return fmt.Errorf("--shell reguires an argument.")
			}
			shellPath = args[1]
			args = args[1:]
		} else if strings.HasPrefix(arg, "--shell=") {
			shellPath = strings.Split(arg, "=")[1]
		} else if arg == "--rsync" {
			if len(args) < 2 {
				return fmt.Errorf("--rsync reguires an argument.")
			}
			rsyncArg = args[1]
			args = args[1:]
		} else if strings.HasPrefix(arg, "--rsync=") {
			rsyncArg = strings.Split(arg, "=")[1]
		} else {
			break
		}

		args = args[1:]
	}

	if zshCompletinFlag {
		fmt.Print(ZSH_COMPLETION)
		return nil
	}

	if bashCompletinFlag {
		fmt.Print(BASH_COMPLETION)
		return nil
	}

	if configFlag {
		ShellExec("$EDITOR " + PerUserConfigFile)
		return nil
	}

	if systemConfigFlag {
		ShellExec("$EDITOR " + SystemWideConfigFile)
		return nil
	}

	var shellContent []byte

	if shellPath != "" {
		if strings.HasPrefix(shellPath, "http://") || strings.HasPrefix(shellPath, "https://") {
			// get script from remote using http.
			if debugFlag {
				fmt.Printf("[zssh debug] get script using http from '%s'\n", shellPath)
			}

			var httpClient *http.Client
			if strings.HasPrefix(shellPath, "https://") {
				tr := &http.Transport{
					TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
				}
				httpClient = &http.Client{Transport: tr}
			} else {
				httpClient = &http.Client{}
			}

			resp, err := httpClient.Get(shellPath)
			if err != nil {
				return err
			}
			defer resp.Body.Close()
			b, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				return err
			}

			shellContent = b
		} else {
			// get script from the file system.
			b, err := ioutil.ReadFile(shellPath)
			if err != nil {
				return err
			}
			shellContent = b
		}
	}

	// set up the lua state.
	L := lua.NewState()
	defer L.Close()

	// load lua custom functions
	LoadFunctions(L)

	if debugFlag {
		fmt.Printf("[zssh debug] loaded lua functions\n")
	}

	// load specific config file
	if configFile != "" {
		_, err := os.Stat(configFile)
		if err != nil {
			return err
		}

		if err := L.DoFile(configFile); err != nil {
			return err
		}

		if debugFlag {
			fmt.Printf("[zssh debug] loaded config file: %s \n", configFile)
		}

	} else {
		// load system wide config
		if _, err := os.Stat(SystemWideConfigFile); err == nil {
			if err := L.DoFile(SystemWideConfigFile); err != nil {
				return err
			}

			if debugFlag {
				fmt.Printf("[zssh debug] loaded config file: %s \n", SystemWideConfigFile)
			}

		}

		// load per-user wide config
		if _, err := os.Stat(PerUserConfigFile); err == nil {
			if err := L.DoFile(PerUserConfigFile); err != nil {
				return err
			}

			if debugFlag {
				fmt.Printf("[zssh debug] loaded config file: %s \n", PerUserConfigFile)
			}

		}
	}

	// generate ssh hosts config
	content, err := GenHostsConfig()
	if err != nil {
		return err
	}

	// only print generated config
	if printFlag {
		fmt.Println(string(content))
		return nil
	}

	// only print hosts list
	if hostsFlag {
		var hosts []*Host
		if len(filters) > 0 {
			hosts = HostsByTags(filters)
		} else {
			hosts = Hosts
		}

		for _, host := range hosts {
			if !host.Hidden {
				if verboseFlag {
					fmt.Printf("%s\t%s\n", host.Name, host.Description)
				} else {
					fmt.Printf("%s\n", host.Name)
				}
			}
		}

		return nil
	}

	if tagsFlag {
		for _, tag := range Tags() {
			fmt.Printf("%s\n", tag)
		}

		return nil
	}

	// generate temporary ssh config file
	tmpFile, err := ioutil.TempFile("", "zssh.ssh_config.")
	if err != nil {
		return err
	}
	defer func() {
		tmpFile.Close()
		os.Remove(tmpFile.Name())

		if debugFlag {
			fmt.Printf("[zssh debug] deleted config file: %s \n", tmpFile.Name())
		}

	}()
	generatedSSHConfigFile := tmpFile.Name()

	if debugFlag {
		fmt.Printf("[zssh debug] generated config file: %s \n", generatedSSHConfigFile)
	}

	// update temporary sss config file
	err = ioutil.WriteFile(generatedSSHConfigFile, content, 0644)
	if err != nil {
		return err
	}

	if shellPath == "" && rsyncArg == "" {
		// get hooks
		var hooks map[string]func() error

		// Limitation!: hooks fires only when the hostname is specified by the first argument.
		if len(args) > 0 {
			hostname := args[0]
			if host := GetHost(hostname); host != nil {
				hooks = host.Hooks
			}
		}

		// run before hook
		if before := hooks["before"]; before != nil {
			if debugFlag {
				fmt.Printf("[zssh debug] run before hook\n")
			}
			err := before()
			if err != nil {
				return err
			}
		}

		// register after hook
		defer func() {
			// after hook
			if after := hooks["after"]; after != nil {
				if debugFlag {
					fmt.Printf("[zssh debug] run after hook\n")
				}
				err := after()
				if err != nil {
					panic(err)
				}
			}
		}()
	}

	// setup ssh command args
	sshComandArgs := []string{"-F", generatedSSHConfigFile}
	sshComandArgs = append(sshComandArgs, args[:]...)
	if shellPath != "" {
		sshComandArgs = append(sshComandArgs, "bash", "-se")
	}

	if rsyncArg != "" {
		// rsync mode.
		if debugFlag {
			fmt.Printf("[zssh debug] use rsync mode.\n")
		}

		rsyncSSHOption := `-e "ssh ` + strings.Join(sshComandArgs, " ") + `"`
		rsyncCommand := "rsync " + rsyncSSHOption + " " + rsyncArg

		if debugFlag {
			fmt.Printf("[zssh debug] real rsync command: %v\n", rsyncCommand)
		}

		err := ShellExec(rsyncCommand)
		if err != nil {
			return err
		}
	} else {
		// normal ssh mode.

		// execute ssh commmand
		cmd := exec.Command("ssh", sshComandArgs[:]...)
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr

		if shellPath != "" {
			// input shell script from stdin.
			cmd.Stdin = bytes.NewBuffer(shellContent)

			if debugFlag {
				fmt.Printf("[zssh debug] shell content:\n%s\n", string(shellContent))
			}

		} else {
			cmd.Stdin = os.Stdin
		}

		if debugFlag {
			fmt.Printf("[zssh debug] real ssh command: %v \n", cmd.Args)
		}

		err = cmd.Run()
		if err != nil {
			return err
		}
	}

	return nil
}

func printUsage() {
	// print usage.
	fmt.Println(`Usage: zssh [<options>] [<ssh options and args...>]

zssh is an extended ssh command.
version ` + Version + `

zssh options:
  --print                 Print generated ssh config.
  --config                Edit per-user config file.
  --system-config         Edit system wide config file.
  --config-file <FILE>    Load configuration from the specific file.
  --hosts                 List hosts. This option can use with additional options.
  --filter <TAG>          (Using with --hosts option) Show only the hosts configured with a tag.
  --verbose               (Using with --hosts option) List hosts with description.
  --tags                  List tags.
  --zsh-completion        Output zsh completion code.
  --debug                 Output debug log

zssh options for convenient functionalities:
  --shell <PATH> <HOSTNAME>              Execute a shell script of the path on the remote host.
  --rsync <rsync options and args...>    Execute rsync using zssh config.

And the following is original ssh command usage...
`)
	// show ssh help
	cmd := exec.Command("ssh")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	cmd.Run()
}

func init() {
	if SystemWideConfigFile == "" {
		SystemWideConfigFile = "/etc/zssh/config.lua"
	}
	if PerUserConfigFile == "" {
		home := userHomeDir()
		PerUserConfigFile = filepath.Join(home, ".zssh/config.lua")
	}
}

var ZSH_COMPLETION = `
_zssh_hosts() {
    local -a __zssh_hosts
    PRE_IFS=$IFS
    IFS=$'\n'
    __zssh_hosts=($(zssh --hosts --verbose | awk -F'\t' '{print $1":"$2}'))
    IFS=$PRE_IFS
    _describe -t host "host" __zssh_hosts
}

_zssh () {
    local curcontext="$curcontext" state line
    typeset -A opt_args

    _arguments \
        '1: :->command'

    case $state in
        command)
            _zssh_hosts
            ;;
        *)
            _files
            ;;
    esac
}

compdef _zssh zssh

`

var BASH_COMPLETION = `
_zssh_hosts() {

}

_zssh () {

}

complete -F _zssh zssh

`

<|endoftext|>


<commit_before>package megos

import (
	"fmt"
	"net/http"
	"net/url"
	"testing"
)

func TestGetURLForStateFile(t *testing.T) {
	setup()
	defer teardown()

	masterNode := client.Master[1]
	u := client.GetURLForStateFile(*masterNode)

	ms := masterNode.String() + "/master/state.json"
	if us := u.String(); ms != us {
		t.Errorf("URLs not the same. Expected %s, got %s", ms, us)
	}
}

func TestGetURLForStateFilePid(t *testing.T) {
	setup()
	defer teardown()

	p := Pid{
		Role: "master",
		Host: "192.168.1.6",
		Port: 5050,
	}

	u := client.GetURLForStateFilePid(p)

	ps := fmt.Sprintf("http://%s:%d/master/state.json", p.Host, p.Port)
	if us := u.String(); ps != us {
		t.Errorf("URLs not the same. Expected %s, got %s", ps, us)
	}
}

func TestGetStateFromCluster_NoNodesOnline(t *testing.T) {
	setup()
	defer teardown()

	u, _ := url.Parse("http://not-existing.example.org/")
	client.Master = []*url.URL{u, u, u}
	state, err := client.GetStateFromCluster()

	if state != nil {
		t.Errorf("State is not nil. Expected nil. Got %+v", state)
	}

	if err == nil {
		t.Error("Error is nil. Expected an error.")
	}
}

func TestGetStateFromCluster(t *testing.T) {
	setup()
	defer teardown()

	mux1.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		c := getContentOfFile("tests/master1.state.json")
		fmt.Fprint(w, string(c))
	})

	state, err := client.GetStateFromCluster()

	if state == nil {
		t.Error("State is nil. Expected valid state struct")
	}

	// Check if some random samples are matching with dummy content in master1.state.json
	if state != nil && (state.Cluster != "chronos1" || state.Version != "0.22.1" || state.Flags.LogDir != "/var/log/mesos" || state.Flags.ZKSessionTimeout != "10secs") {
		t.Error("Random samples are not matching with test mock data.")
	}

	if err != nil {
		t.Errorf("Error is not nil. Expected nil, got %s.", err)
	}
}

// Missing tests:
// * GetStateFromLeader
// * GetStateFromPid
<commit_msg>Added unit tests for GetStateFromLeader<commit_after>package megos

import (
	"fmt"
	"net/http"
	"net/url"
	"testing"
)

func TestGetURLForStateFile(t *testing.T) {
	setup()
	defer teardown()

	masterNode := client.Master[1]
	u := client.GetURLForStateFile(*masterNode)

	ms := masterNode.String() + "/master/state.json"
	if us := u.String(); ms != us {
		t.Errorf("URLs not the same. Expected %s, got %s", ms, us)
	}
}

func TestGetURLForStateFilePid(t *testing.T) {
	setup()
	defer teardown()

	p := Pid{
		Role: "master",
		Host: "192.168.1.6",
		Port: 5050,
	}

	u := client.GetURLForStateFilePid(p)

	ps := fmt.Sprintf("http://%s:%d/master/state.json", p.Host, p.Port)
	if us := u.String(); ps != us {
		t.Errorf("URLs not the same. Expected %s, got %s", ps, us)
	}
}

func TestGetStateFromCluster_NoNodesOnline(t *testing.T) {
	setup()
	defer teardown()

	u, _ := url.Parse("http://not-existing.example.org/")
	client.Master = []*url.URL{u, u, u}
	state, err := client.GetStateFromCluster()

	if state != nil {
		t.Errorf("State is not nil. Expected nil. Got %+v", state)
	}

	if err == nil {
		t.Error("Error is nil. Expected an error.")
	}
}

func TestGetStateFromCluster(t *testing.T) {
	setup()
	defer teardown()

	mux1.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		c := getContentOfFile("tests/master1.state.json")
		fmt.Fprint(w, string(c))
	})

	state, err := client.GetStateFromCluster()

	if state == nil {
		t.Error("State is nil. Expected valid state struct")
	}

	// Check if some random samples are matching with dummy content in master1.state.json
	if state != nil && (state.Cluster != "chronos1" || state.Version != "0.22.1" || state.Flags.LogDir != "/var/log/mesos" || state.Flags.ZKSessionTimeout != "10secs") {
		t.Error("Random samples are not matching with test mock data.")
	}

	if err != nil {
		t.Errorf("Error is not nil. Expected nil, got %s.", err)
	}
}

func TestGetStateFromLeader_NoLeader(t *testing.T) {
	setup()
	defer teardown()

	state, err := client.GetStateFromLeader()

	if state != nil {
		t.Errorf("State is not nil. Expected nil. Got %+v", state)
	}

	if err == nil {
		t.Error("Error is nil. Expected an error.")
	}
}

func TestGetStateFromLeader_NoLeaderOnline(t *testing.T) {
	setup()
	defer teardown()

	client.Leader = &Pid{
		Role: "master",
		Host: "not-existing.example.org",
	}
	state, err := client.GetStateFromLeader()

	if state != nil {
		t.Errorf("State is not nil. Expected nil. Got %+v", state)
	}

	if err == nil {
		t.Error("Error is nil. Expected an error.")
	}
}

func GetStateFromLeader(t *testing.T) {
	setup()
	defer teardown()

	mux1.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		testMethod(t, r, "GET")
		c := getContentOfFile("tests/master1.state.json")
		fmt.Fprint(w, string(c))
	})

	state, err := client.GetStateFromLeader()

	if state == nil {
		t.Error("State is nil. Expected valid state struct")
	}

	// Check if some random samples are matching with dummy content in master1.state.json
	if state != nil && (state.Cluster != "chronos1" || state.Version != "0.22.1" || state.Flags.LogDir != "/var/log/mesos" || state.Flags.ZKSessionTimeout != "10secs") {
		t.Error("Random samples are not matching with test mock data.")
	}

	if err != nil {
		t.Errorf("Error is not nil. Expected nil, got %s.", err)
	}
}

// Missing tests:
// * GetStateFromPid
<|endoftext|>


<commit_before>// +build integration

package tests

import (
	"testing"

	"github.com/deis/deis/tests/utils"
)

var (
	configListCmd  = "config:list --app={{.AppName}}"
	configSetCmd   = "config:set FOO=讲台 --app={{.AppName}}"
	configSet2Cmd  = "config:set FOO=10 --app={{.AppName}}"
	configUnsetCmd = "config:unset FOO --app={{.AppName}}"
)

func TestConfig(t *testing.T) {
	params := configSetup(t)
	configSetTest(t, params)
	configListTest(t, params, false)
	appsOpenTest(t, params)
	configUnsetTest(t, params)
	configListTest(t, params, true)
	limitsSetTest(t, params, 4)
	appsOpenTest(t, params)
	limitsUnsetTest(t, params, 6)
	appsOpenTest(t, params)
	//tagsTest(t, params, 8)
	appsOpenTest(t, params)
	utils.AppsDestroyTest(t, params)
}

func configSetup(t *testing.T) *utils.DeisTestConfig {
	cfg := utils.GetGlobalConfig()
	cfg.AppName = "configsample"
	utils.Execute(t, authLoginCmd, cfg, false, "")
	utils.Execute(t, gitCloneCmd, cfg, false, "")
	if err := utils.Chdir(cfg.ExampleApp); err != nil {
		t.Fatal(err)
	}
	utils.Execute(t, appsCreateCmd, cfg, false, "")
	utils.Execute(t, gitPushCmd, cfg, false, "")
	if err := utils.Chdir(".."); err != nil {
		t.Fatal(err)
	}
	return cfg
}

func configListTest(
	t *testing.T, params *utils.DeisTestConfig, notflag bool) {
	utils.CheckList(t, configListCmd, params, "FOO", notflag)
}

func configSetTest(t *testing.T, params *utils.DeisTestConfig) {
	utils.Execute(t, configSetCmd, params, false, "讲台")
	utils.CheckList(t, appsInfoCmd, params, "(v3)", false)
	utils.Execute(t, configSet2Cmd, params, false, "10")
	utils.CheckList(t, appsInfoCmd, params, "(v4)", false)
}

func configUnsetTest(t *testing.T, params *utils.DeisTestConfig) {
	utils.Execute(t, configUnsetCmd, params, false, "")
	utils.CheckList(t, appsInfoCmd, params, "(v5)", false)
}
<commit_msg>test(tests): add config:unset tests<commit_after>// +build integration

package tests

import (
	"testing"

	"github.com/deis/deis/tests/utils"
)

var (
	configListCmd  = "config:list --app={{.AppName}}"
	configSetCmd   = "config:set FOO=讲台 --app={{.AppName}}"
	configSet2Cmd  = "config:set FOO=10 --app={{.AppName}}"
	configUnsetCmd = "config:unset FOO --app={{.AppName}}"
)

func TestConfig(t *testing.T) {
	params := configSetup(t)
	configSetTest(t, params)
	configListTest(t, params, false)
	appsOpenTest(t, params)
	configUnsetTest(t, params)
	configListTest(t, params, true)
	limitsSetTest(t, params, 4)
	appsOpenTest(t, params)
	limitsUnsetTest(t, params, 6)
	appsOpenTest(t, params)
	//tagsTest(t, params, 8)
	appsOpenTest(t, params)
	utils.AppsDestroyTest(t, params)
}

func configSetup(t *testing.T) *utils.DeisTestConfig {
	cfg := utils.GetGlobalConfig()
	cfg.AppName = "configsample"
	utils.Execute(t, authLoginCmd, cfg, false, "")
	utils.Execute(t, gitCloneCmd, cfg, false, "")
	if err := utils.Chdir(cfg.ExampleApp); err != nil {
		t.Fatal(err)
	}
	utils.Execute(t, appsCreateCmd, cfg, false, "")
	utils.Execute(t, gitPushCmd, cfg, false, "")
	if err := utils.Chdir(".."); err != nil {
		t.Fatal(err)
	}
	return cfg
}

func configListTest(
	t *testing.T, params *utils.DeisTestConfig, notflag bool) {
	utils.CheckList(t, configListCmd, params, "FOO", notflag)
}

func configSetTest(t *testing.T, params *utils.DeisTestConfig) {
	utils.Execute(t, configSetCmd, params, false, "讲台")
	utils.CheckList(t, appsInfoCmd, params, "(v3)", false)
	utils.Execute(t, configSet2Cmd, params, false, "10")
	utils.CheckList(t, appsInfoCmd, params, "(v4)", false)
}

func configUnsetTest(t *testing.T, params *utils.DeisTestConfig) {
	utils.Execute(t, configUnsetCmd, params, false, "")
	utils.CheckList(t, appsInfoCmd, params, "(v5)", false)
	utils.CheckList(t, "run env --app={{.AppName}}", params, "FOO", true)
}
<|endoftext|>