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

// Garbage collector: sweeping

package runtime

import "unsafe"

var sweep sweepdata

// State of background sweep.
type sweepdata struct {
	lock    mutex
	g       *g
	parked  bool
	started bool

	spanidx uint32 // background sweeper position

	nbgsweep    uint32
	npausesweep uint32
}

//go:nowritebarrier
func finishsweep_m() {
	// The world is stopped so we should be able to complete the sweeps
	// quickly.
	for sweepone() != ^uintptr(0) {
		sweep.npausesweep++
	}

	// There may be some other spans being swept concurrently that
	// we need to wait for. If finishsweep_m is done with the world stopped
	// this code is not required.
	sg := mheap_.sweepgen
	for _, s := range work.spans {
		if s.sweepgen != sg && s.state == _MSpanInUse {
			mSpan_EnsureSwept(s)
		}
	}
}

func bgsweep(c chan int) {
	sweep.g = getg()

	lock(&sweep.lock)
	sweep.parked = true
	c <- 1
	goparkunlock(&sweep.lock, "GC sweep wait", traceEvGoBlock, 1)

	for {
		for gosweepone() != ^uintptr(0) {
			sweep.nbgsweep++
			Gosched()
		}
		lock(&sweep.lock)
		if !gosweepdone() {
			// This can happen if a GC runs between
			// gosweepone returning ^0 above
			// and the lock being acquired.
			unlock(&sweep.lock)
			continue
		}
		sweep.parked = true
		goparkunlock(&sweep.lock, "GC sweep wait", traceEvGoBlock, 1)
	}
}

// sweeps one span
// returns number of pages returned to heap, or ^uintptr(0) if there is nothing to sweep
//go:nowritebarrier
func sweepone() uintptr {
	_g_ := getg()

	// increment locks to ensure that the goroutine is not preempted
	// in the middle of sweep thus leaving the span in an inconsistent state for next GC
	_g_.m.locks++
	sg := mheap_.sweepgen
	for {
		idx := xadd(&sweep.spanidx, 1) - 1
		if idx >= uint32(len(work.spans)) {
			mheap_.sweepdone = 1
			_g_.m.locks--
			return ^uintptr(0)
		}
		s := work.spans[idx]
		if s.state != mSpanInUse {
			s.sweepgen = sg
			continue
		}
		if s.sweepgen != sg-2 || !cas(&s.sweepgen, sg-2, sg-1) {
			continue
		}
		npages := s.npages
		if !mSpan_Sweep(s, false) {
			npages = 0
		}
		_g_.m.locks--
		return npages
	}
}

//go:nowritebarrier
func gosweepone() uintptr {
	var ret uintptr
	systemstack(func() {
		ret = sweepone()
	})
	return ret
}

//go:nowritebarrier
func gosweepdone() bool {
	return mheap_.sweepdone != 0
}

// Returns only when span s has been swept.
//go:nowritebarrier
func mSpan_EnsureSwept(s *mspan) {
	// Caller must disable preemption.
	// Otherwise when this function returns the span can become unswept again
	// (if GC is triggered on another goroutine).
	_g_ := getg()
	if _g_.m.locks == 0 && _g_.m.mallocing == 0 && _g_ != _g_.m.g0 {
		throw("MSpan_EnsureSwept: m is not locked")
	}

	sg := mheap_.sweepgen
	if atomicload(&s.sweepgen) == sg {
		return
	}
	// The caller must be sure that the span is a MSpanInUse span.
	if cas(&s.sweepgen, sg-2, sg-1) {
		mSpan_Sweep(s, false)
		return
	}
	// unfortunate condition, and we don't have efficient means to wait
	for atomicload(&s.sweepgen) != sg {
		osyield()
	}
}

// Sweep frees or collects finalizers for blocks not marked in the mark phase.
// It clears the mark bits in preparation for the next GC round.
// Returns true if the span was returned to heap.
// If preserve=true, don't return it to heap nor relink in MCentral lists;
// caller takes care of it.
//TODO go:nowritebarrier
func mSpan_Sweep(s *mspan, preserve bool) bool {
	// It's critical that we enter this function with preemption disabled,
	// GC must not start while we are in the middle of this function.
	_g_ := getg()
	if _g_.m.locks == 0 && _g_.m.mallocing == 0 && _g_ != _g_.m.g0 {
		throw("MSpan_Sweep: m is not locked")
	}
	sweepgen := mheap_.sweepgen
	if s.state != mSpanInUse || s.sweepgen != sweepgen-1 {
		print("MSpan_Sweep: state=", s.state, " sweepgen=", s.sweepgen, " mheap.sweepgen=", sweepgen, "\n")
		throw("MSpan_Sweep: bad span state")
	}

	if trace.enabled {
		traceGCSweepStart()
	}

	cl := s.sizeclass
	size := s.elemsize
	res := false
	nfree := 0

	var head, end gclinkptr

	c := _g_.m.mcache
	sweepgenset := false

	// Mark any free objects in this span so we don't collect them.
	for link := s.freelist; link.ptr() != nil; link = link.ptr().next {
		heapBitsForAddr(uintptr(link)).setMarkedNonAtomic()
	}

	// Unlink & free special records for any objects we're about to free.
	specialp := &s.specials
	special := *specialp
	for special != nil {
		// A finalizer can be set for an inner byte of an object, find object beginning.
		p := uintptr(s.start<<_PageShift) + uintptr(special.offset)/size*size
		hbits := heapBitsForAddr(p)
		if !hbits.isMarked() {
			// Find the exact byte for which the special was setup
			// (as opposed to object beginning).
			p := uintptr(s.start<<_PageShift) + uintptr(special.offset)
			// about to free object: splice out special record
			y := special
			special = special.next
			*specialp = special
			if !freespecial(y, unsafe.Pointer(p), size, false) {
				// stop freeing of object if it has a finalizer
				hbits.setMarkedNonAtomic()
			}
		} else {
			// object is still live: keep special record
			specialp = &special.next
			special = *specialp
		}
	}

	// Sweep through n objects of given size starting at p.
	// This thread owns the span now, so it can manipulate
	// the block bitmap without atomic operations.

	size, n, _ := s.layout()
	heapBitsSweepSpan(s.base(), size, n, func(p uintptr) {
		// At this point we know that we are looking at garbage object
		// that needs to be collected.
		if debug.allocfreetrace != 0 {
			tracefree(unsafe.Pointer(p), size)
		}

		// Reset to allocated+noscan.
		if cl == 0 {
			// Free large span.
			if preserve {
				throw("can't preserve large span")
			}
			heapBitsForSpan(p).initSpan(s.layout())
			s.needzero = 1

			// important to set sweepgen before returning it to heap
			atomicstore(&s.sweepgen, sweepgen)
			sweepgenset = true

			// NOTE(rsc,dvyukov): The original implementation of efence
			// in CL 22060046 used SysFree instead of SysFault, so that
			// the operating system would eventually give the memory
			// back to us again, so that an efence program could run
			// longer without running out of memory. Unfortunately,
			// calling SysFree here without any kind of adjustment of the
			// heap data structures means that when the memory does
			// come back to us, we have the wrong metadata for it, either in
			// the MSpan structures or in the garbage collection bitmap.
			// Using SysFault here means that the program will run out of
			// memory fairly quickly in efence mode, but at least it won't
			// have mysterious crashes due to confused memory reuse.
			// It should be possible to switch back to SysFree if we also
			// implement and then call some kind of MHeap_DeleteSpan.
			if debug.efence > 0 {
				s.limit = 0 // prevent mlookup from finding this span
				sysFault(unsafe.Pointer(p), size)
			} else {
				mHeap_Free(&mheap_, s, 1)
			}
			c.local_nlargefree++
			c.local_largefree += size
			res = true
		} else {
			// Free small object.
			if size > 2*ptrSize {
				*(*uintptr)(unsafe.Pointer(p + ptrSize)) = uintptrMask & 0xdeaddeaddeaddead // mark as "needs to be zeroed"
			} else if size > ptrSize {
				*(*uintptr)(unsafe.Pointer(p + ptrSize)) = 0
			}
			if head.ptr() == nil {
				head = gclinkptr(p)
			} else {
				end.ptr().next = gclinkptr(p)
			}
			end = gclinkptr(p)
			end.ptr().next = gclinkptr(0x0bade5)
			nfree++
		}
	})

	// We need to set s.sweepgen = h.sweepgen only when all blocks are swept,
	// because of the potential for a concurrent free/SetFinalizer.
	// But we need to set it before we make the span available for allocation
	// (return it to heap or mcentral), because allocation code assumes that a
	// span is already swept if available for allocation.
	if !sweepgenset && nfree == 0 {
		// The span must be in our exclusive ownership until we update sweepgen,
		// check for potential races.
		if s.state != mSpanInUse || s.sweepgen != sweepgen-1 {
			print("MSpan_Sweep: state=", s.state, " sweepgen=", s.sweepgen, " mheap.sweepgen=", sweepgen, "\n")
			throw("MSpan_Sweep: bad span state after sweep")
		}
		atomicstore(&s.sweepgen, sweepgen)
	}
	if nfree > 0 {
		c.local_nsmallfree[cl] += uintptr(nfree)
		res = mCentral_FreeSpan(&mheap_.central[cl].mcentral, s, int32(nfree), head, end, preserve)
		// MCentral_FreeSpan updates sweepgen
	}
	if trace.enabled {
		traceGCSweepDone()
		traceNextGC()
	}
	return res
}
<commit_msg>runtime: remove unnecessary traceNextGC<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Garbage collector: sweeping

package runtime

import "unsafe"

var sweep sweepdata

// State of background sweep.
type sweepdata struct {
	lock    mutex
	g       *g
	parked  bool
	started bool

	spanidx uint32 // background sweeper position

	nbgsweep    uint32
	npausesweep uint32
}

//go:nowritebarrier
func finishsweep_m() {
	// The world is stopped so we should be able to complete the sweeps
	// quickly.
	for sweepone() != ^uintptr(0) {
		sweep.npausesweep++
	}

	// There may be some other spans being swept concurrently that
	// we need to wait for. If finishsweep_m is done with the world stopped
	// this code is not required.
	sg := mheap_.sweepgen
	for _, s := range work.spans {
		if s.sweepgen != sg && s.state == _MSpanInUse {
			mSpan_EnsureSwept(s)
		}
	}
}

func bgsweep(c chan int) {
	sweep.g = getg()

	lock(&sweep.lock)
	sweep.parked = true
	c <- 1
	goparkunlock(&sweep.lock, "GC sweep wait", traceEvGoBlock, 1)

	for {
		for gosweepone() != ^uintptr(0) {
			sweep.nbgsweep++
			Gosched()
		}
		lock(&sweep.lock)
		if !gosweepdone() {
			// This can happen if a GC runs between
			// gosweepone returning ^0 above
			// and the lock being acquired.
			unlock(&sweep.lock)
			continue
		}
		sweep.parked = true
		goparkunlock(&sweep.lock, "GC sweep wait", traceEvGoBlock, 1)
	}
}

// sweeps one span
// returns number of pages returned to heap, or ^uintptr(0) if there is nothing to sweep
//go:nowritebarrier
func sweepone() uintptr {
	_g_ := getg()

	// increment locks to ensure that the goroutine is not preempted
	// in the middle of sweep thus leaving the span in an inconsistent state for next GC
	_g_.m.locks++
	sg := mheap_.sweepgen
	for {
		idx := xadd(&sweep.spanidx, 1) - 1
		if idx >= uint32(len(work.spans)) {
			mheap_.sweepdone = 1
			_g_.m.locks--
			return ^uintptr(0)
		}
		s := work.spans[idx]
		if s.state != mSpanInUse {
			s.sweepgen = sg
			continue
		}
		if s.sweepgen != sg-2 || !cas(&s.sweepgen, sg-2, sg-1) {
			continue
		}
		npages := s.npages
		if !mSpan_Sweep(s, false) {
			npages = 0
		}
		_g_.m.locks--
		return npages
	}
}

//go:nowritebarrier
func gosweepone() uintptr {
	var ret uintptr
	systemstack(func() {
		ret = sweepone()
	})
	return ret
}

//go:nowritebarrier
func gosweepdone() bool {
	return mheap_.sweepdone != 0
}

// Returns only when span s has been swept.
//go:nowritebarrier
func mSpan_EnsureSwept(s *mspan) {
	// Caller must disable preemption.
	// Otherwise when this function returns the span can become unswept again
	// (if GC is triggered on another goroutine).
	_g_ := getg()
	if _g_.m.locks == 0 && _g_.m.mallocing == 0 && _g_ != _g_.m.g0 {
		throw("MSpan_EnsureSwept: m is not locked")
	}

	sg := mheap_.sweepgen
	if atomicload(&s.sweepgen) == sg {
		return
	}
	// The caller must be sure that the span is a MSpanInUse span.
	if cas(&s.sweepgen, sg-2, sg-1) {
		mSpan_Sweep(s, false)
		return
	}
	// unfortunate condition, and we don't have efficient means to wait
	for atomicload(&s.sweepgen) != sg {
		osyield()
	}
}

// Sweep frees or collects finalizers for blocks not marked in the mark phase.
// It clears the mark bits in preparation for the next GC round.
// Returns true if the span was returned to heap.
// If preserve=true, don't return it to heap nor relink in MCentral lists;
// caller takes care of it.
//TODO go:nowritebarrier
func mSpan_Sweep(s *mspan, preserve bool) bool {
	// It's critical that we enter this function with preemption disabled,
	// GC must not start while we are in the middle of this function.
	_g_ := getg()
	if _g_.m.locks == 0 && _g_.m.mallocing == 0 && _g_ != _g_.m.g0 {
		throw("MSpan_Sweep: m is not locked")
	}
	sweepgen := mheap_.sweepgen
	if s.state != mSpanInUse || s.sweepgen != sweepgen-1 {
		print("MSpan_Sweep: state=", s.state, " sweepgen=", s.sweepgen, " mheap.sweepgen=", sweepgen, "\n")
		throw("MSpan_Sweep: bad span state")
	}

	if trace.enabled {
		traceGCSweepStart()
	}

	cl := s.sizeclass
	size := s.elemsize
	res := false
	nfree := 0

	var head, end gclinkptr

	c := _g_.m.mcache
	sweepgenset := false

	// Mark any free objects in this span so we don't collect them.
	for link := s.freelist; link.ptr() != nil; link = link.ptr().next {
		heapBitsForAddr(uintptr(link)).setMarkedNonAtomic()
	}

	// Unlink & free special records for any objects we're about to free.
	specialp := &s.specials
	special := *specialp
	for special != nil {
		// A finalizer can be set for an inner byte of an object, find object beginning.
		p := uintptr(s.start<<_PageShift) + uintptr(special.offset)/size*size
		hbits := heapBitsForAddr(p)
		if !hbits.isMarked() {
			// Find the exact byte for which the special was setup
			// (as opposed to object beginning).
			p := uintptr(s.start<<_PageShift) + uintptr(special.offset)
			// about to free object: splice out special record
			y := special
			special = special.next
			*specialp = special
			if !freespecial(y, unsafe.Pointer(p), size, false) {
				// stop freeing of object if it has a finalizer
				hbits.setMarkedNonAtomic()
			}
		} else {
			// object is still live: keep special record
			specialp = &special.next
			special = *specialp
		}
	}

	// Sweep through n objects of given size starting at p.
	// This thread owns the span now, so it can manipulate
	// the block bitmap without atomic operations.

	size, n, _ := s.layout()
	heapBitsSweepSpan(s.base(), size, n, func(p uintptr) {
		// At this point we know that we are looking at garbage object
		// that needs to be collected.
		if debug.allocfreetrace != 0 {
			tracefree(unsafe.Pointer(p), size)
		}

		// Reset to allocated+noscan.
		if cl == 0 {
			// Free large span.
			if preserve {
				throw("can't preserve large span")
			}
			heapBitsForSpan(p).initSpan(s.layout())
			s.needzero = 1

			// important to set sweepgen before returning it to heap
			atomicstore(&s.sweepgen, sweepgen)
			sweepgenset = true

			// NOTE(rsc,dvyukov): The original implementation of efence
			// in CL 22060046 used SysFree instead of SysFault, so that
			// the operating system would eventually give the memory
			// back to us again, so that an efence program could run
			// longer without running out of memory. Unfortunately,
			// calling SysFree here without any kind of adjustment of the
			// heap data structures means that when the memory does
			// come back to us, we have the wrong metadata for it, either in
			// the MSpan structures or in the garbage collection bitmap.
			// Using SysFault here means that the program will run out of
			// memory fairly quickly in efence mode, but at least it won't
			// have mysterious crashes due to confused memory reuse.
			// It should be possible to switch back to SysFree if we also
			// implement and then call some kind of MHeap_DeleteSpan.
			if debug.efence > 0 {
				s.limit = 0 // prevent mlookup from finding this span
				sysFault(unsafe.Pointer(p), size)
			} else {
				mHeap_Free(&mheap_, s, 1)
			}
			c.local_nlargefree++
			c.local_largefree += size
			res = true
		} else {
			// Free small object.
			if size > 2*ptrSize {
				*(*uintptr)(unsafe.Pointer(p + ptrSize)) = uintptrMask & 0xdeaddeaddeaddead // mark as "needs to be zeroed"
			} else if size > ptrSize {
				*(*uintptr)(unsafe.Pointer(p + ptrSize)) = 0
			}
			if head.ptr() == nil {
				head = gclinkptr(p)
			} else {
				end.ptr().next = gclinkptr(p)
			}
			end = gclinkptr(p)
			end.ptr().next = gclinkptr(0x0bade5)
			nfree++
		}
	})

	// We need to set s.sweepgen = h.sweepgen only when all blocks are swept,
	// because of the potential for a concurrent free/SetFinalizer.
	// But we need to set it before we make the span available for allocation
	// (return it to heap or mcentral), because allocation code assumes that a
	// span is already swept if available for allocation.
	if !sweepgenset && nfree == 0 {
		// The span must be in our exclusive ownership until we update sweepgen,
		// check for potential races.
		if s.state != mSpanInUse || s.sweepgen != sweepgen-1 {
			print("MSpan_Sweep: state=", s.state, " sweepgen=", s.sweepgen, " mheap.sweepgen=", sweepgen, "\n")
			throw("MSpan_Sweep: bad span state after sweep")
		}
		atomicstore(&s.sweepgen, sweepgen)
	}
	if nfree > 0 {
		c.local_nsmallfree[cl] += uintptr(nfree)
		res = mCentral_FreeSpan(&mheap_.central[cl].mcentral, s, int32(nfree), head, end, preserve)
		// MCentral_FreeSpan updates sweepgen
	}
	if trace.enabled {
		traceGCSweepDone()
	}
	return res
}
<|endoftext|>


<commit_before>package scummatlas

import (
	"binaryutils"
	"fmt"
	"image"
	"image/color"
	"io/ioutil"
	"log"
	"os"
)

const debugImage = false

const (
	_ = iota
	METHOD_UNKNOWN
	METHOD_UNCOMPRESSED
	METHOD_ONE
	METHOD_TWO
	HORIZONTAL
	VERTICAL
	TRANSP
	NO_TRANSP
)

func parsePalette(data []byte) color.Palette {
	var palette = make(color.Palette, 0, 256)
	for i := 0; i < len(data); i += 3 {
		color := color.RGBA{
			data[i],
			data[i+1],
			data[i+2],
			255,
		}
		palette = append(palette, color)
	}

	return palette
}

func parseImage(data []byte, zBuffers int, width int, height int, pal color.Palette, transpIndex uint8) *image.RGBA {
	// DISABLE LOGGING
	log.SetOutput(os.Stdout)
	log.SetOutput(ioutil.Discard)

	if string(data[8:12]) != "SMAP" {
		panic("No stripe table found")
	}

	stripeCount := width / 8
	log.Println("SmapSize", BE32(data, 12))

	log.Println("There should be ", stripeCount, "stripes")
	offsets := make([]int, 0, stripeCount)
	stripeOffset := 0
	for i := 0; i < stripeCount; i++ {
		stripeOffset = LE32(data, 16+4*i)
		offsets = append(offsets, stripeOffset)
	}

	img := image.NewRGBA(image.Rect(0, 0, width, height))

	log.Println("Decoding image")
	log.Println("Stripes information")
	log.Println("\n#ID\tCode\tMethod\tDirect\tPalInSz\tTransp")
	for i := 0; i < stripeCount; i++ {
		offset := offsets[i]
		size := len(data) - offset
		if i < stripeCount-1 {
			size = offsets[i+1] - offsets[i]
		}
		drawStripe(img, i, data[8+offset:8+offset+size], pal, transpIndex)
	}
	log.Println("image decoded")
	log.SetOutput(os.Stdout)
	return img
}

func drawStripe(img *image.RGBA, stripNumber int, data []byte, pal color.Palette, transpIndex uint8) {

	method, direction, transparent, paletteLength := getCompressionMethod(stripNumber, data[0])

	height := img.Rect.Size().Y
	totalPixels := 8 * height
	currentPixel := 0

	curPal := uint8(data[1])
	curSubs := uint8(1)

	setColor := func() {
		var x, y int
		if direction == HORIZONTAL {
			x = currentPixel % 8
			y = (currentPixel + 1) / 8
		} else {
			y = currentPixel % height
			x = currentPixel / height
		}
		if x == 7 {
			x = -1
		}
		x += 8 * stripNumber
		if curPal >= 0 {
			if transparent == TRANSP && curPal == transpIndex {
				log.Println("TRANSPARENT")
				img.Set(x, y, color.RGBA{0, 0, 0, 0})
			} else {
				img.Set(x, y, pal[curPal])
			}
		} else {
			panic("Out of palette")
		}
		currentPixel++
	}

	bs := binaryutils.NewBitStream(data[2:])

	if method == METHOD_TWO {
		for currentPixel < totalPixels {
			if bs.GetBit() == 1 {
				if bs.GetBit() == 1 {
					palShift := bs.GetBits(3) - 4
					if palShift == 0 {
						length := int(bs.GetBits(8))
						for i := 0; i < length; i++ {
							setColor()
						}
					} else {
						curPal += palShift
						setColor()
					}
				} else { // 10 Read new palette index
					curPal = bs.GetBits(paletteLength)
					setColor()
				}
			} else { // 0 Draw next pixel with current palette index
				setColor()
			}
		}
	} else {
		for currentPixel < totalPixels {
			if bs.GetBit() == 1 {
				if bs.GetBit() == 1 {
					if bs.GetBit() == 1 {
						//Negate the subtraction variable. Subtract it from the palette index, and draw the next pixel.
						curSubs = -curSubs
					}
					//Subtract the subtraction variable from the palette index, and draw the next pixel.
					curPal -= curSubs
				} else { //10, read new palette index
					curPal = bs.GetBits(paletteLength)
				}
			} else { // 0, draw next pixel with current palette index
			}
			setColor()
		}
	}
}

func getCompressionMethod(stripNumber int, code byte) (
	method int,
	direction int,
	transparent int,
	paletteLength uint8) {
	var substraction byte

	direction = HORIZONTAL
	method = METHOD_UNKNOWN
	transparent = NO_TRANSP
	substraction = 0x00
	if code == 0x01 {
		method = METHOD_UNCOMPRESSED
	} else if code >= 0x0e && code <= 0x30 {
		method = METHOD_ONE
	} else if code >= 0x40 && code <= 0x80 {
		method = METHOD_TWO
	}

	if (code >= 0x03 && code <= 0x12) ||
		(code >= 0x22 && code <= 0x26) {
		direction = VERTICAL
	}

	if (code >= 0x22 && code <= 0x30) ||
		(code >= 0x54 && code <= 0x6c) {
		transparent = TRANSP
	}

	codes := []uint8{0x0e, 0x18, 0x22, 0x2c, 0x40, 0x54, 0x68, 0x7c}

	for _, v := range codes {
		if code >= v && code <= v+4 {
			substraction = v - 4
			break
		}
	}

	out := fmt.Sprintf("%v\t0x%X\t", stripNumber, code)
	if method == METHOD_ONE {
		out += "   1"
	} else if method == METHOD_TWO {
		out += "   2"
	}
	if direction == VERTICAL {
		out += "\tVert"
	} else {
		out += "\tHoriz"
	}
	out += "\t" + string(code-substraction)
	if transparent == TRANSP {
		out += "\tYes"
	} else {
		out += "\tNo"
	}
	log.Println(out)
	return method, direction, transparent, code - substraction
}
<commit_msg>Clean image stripe<commit_after>package scummatlas

import (
	"binaryutils"
	"fmt"
	"image"
	"image/color"
	"io/ioutil"
	"log"
	"os"
)

const debugImage = false

const (
	_ = iota
	METHOD_UNKNOWN
	METHOD_UNCOMPRESSED
	METHOD_ONE
	METHOD_TWO
	HORIZONTAL
	VERTICAL
	TRANSP
	NO_TRANSP
)

func parsePalette(data []byte) color.Palette {
	var palette = make(color.Palette, 0, 256)
	for i := 0; i < len(data); i += 3 {
		color := color.RGBA{
			data[i],
			data[i+1],
			data[i+2],
			255,
		}
		palette = append(palette, color)
	}

	return palette
}

func parseImage(data []byte, zBuffers int, width int, height int, pal color.Palette, transpIndex uint8) *image.RGBA {
	// DISABLE LOGGING
	log.SetOutput(os.Stdout)
	log.SetOutput(ioutil.Discard)

	if string(data[8:12]) != "SMAP" {
		panic("No stripe table found")
	}

	stripeCount := width / 8
	log.Println("SmapSize", BE32(data, 12))

	log.Println("There should be ", stripeCount, "stripes")
	offsets := make([]int, 0, stripeCount)
	stripeOffset := 0
	for i := 0; i < stripeCount; i++ {
		stripeOffset = LE32(data, 16+4*i)
		offsets = append(offsets, stripeOffset)
	}

	img := image.NewRGBA(image.Rect(0, 0, width, height))

	log.Println("Decoding image")
	log.Println("Stripes information")
	log.Println("\n#ID\tCode\tMethod\tDirect\tPalInSz\tTransp")
	for i := 0; i < stripeCount; i++ {
		offset := offsets[i]
		size := len(data) - offset
		if i < stripeCount-1 {
			size = offsets[i+1] - offsets[i]
		}
		drawStripe(img, i, data[8+offset:8+offset+size], pal, transpIndex)
	}
	log.Println("image decoded")
	log.SetOutput(os.Stdout)
	return img
}

func drawStripe(img *image.RGBA, stripNumber int, data []byte, pal color.Palette, transpIndex uint8) {

	method, direction, transparent, paletteLength := getCompressionMethod(stripNumber, data[0])

	height := img.Rect.Size().Y
	totalPixels := 8 * height
	currentPixel := 0

	curPal := uint8(data[1])
	curSubs := uint8(1)

	setColor := func() {
		var x, y int
		if direction == HORIZONTAL {
			x = currentPixel % 8
			y = (currentPixel + 1) / 8
		} else {
			y = currentPixel % height
			x = currentPixel / height
		}
		if x == 7 {
			x = -1
		}
		x += 8 * stripNumber
		if curPal >= 0 {
			if transparent == TRANSP && curPal == transpIndex {
				log.Println("TRANSPARENT")
				img.Set(x, y, color.RGBA{0, 0, 0, 0})
			} else {
				img.Set(x, y, pal[curPal])
			}
		} else {
			panic("Out of palette")
		}
		currentPixel++
	}

	bs := binaryutils.NewBitStream(data[2:])

	if method == METHOD_TWO {
		for currentPixel < totalPixels {
			if bs.GetBit() == 1 {
				if bs.GetBit() == 1 {
					palShift := bs.GetBits(3) - 4
					if palShift == 0 {
						length := int(bs.GetBits(8))
						for i := 0; i < length; i++ {
							setColor()
						}
					} else {
						curPal += palShift
						setColor()
					}
				} else { // 10 Read new palette index
					curPal = bs.GetBits(paletteLength)
					setColor()
				}
			} else { // 0 Draw next pixel with current palette index
				setColor()
			}
		}
	} else {
		for currentPixel < totalPixels {
			if bs.GetBit() == 1 {
				if bs.GetBit() == 1 {
					if bs.GetBit() == 1 {
						//Negate the subtraction variable. Subtract it from the palette index, and draw the next pixel.
						curSubs = -curSubs
					}
					//Subtract the subtraction variable from the palette index, and draw the next pixel.
					curPal -= curSubs
				} else { //10, read new palette index
					curPal = bs.GetBits(paletteLength)
				}
			} else { // 0, draw next pixel with current palette index
			}
			setColor()
		}
	}
}

func getCompressionMethod(stripNumber int, code byte) (
	method int,
	direction int,
	transparent int,
	paletteLength uint8) {

	direction = HORIZONTAL
	method = METHOD_UNKNOWN
	transparent = NO_TRANSP
	var substraction uint8

	if code == 0x01 {
		method = METHOD_UNCOMPRESSED
	} else if code >= 0x0e && code <= 0x30 {
		method = METHOD_ONE
	} else if code >= 0x40 && code <= 0x80 {
		method = METHOD_TWO
	}

	if (code >= 0x03 && code <= 0x12) ||
		(code >= 0x22 && code <= 0x26) {
		direction = VERTICAL
	}

	if (code >= 0x22 && code <= 0x30) ||
		(code >= 0x54 && code <= 0x6c) {
		transparent = TRANSP
	}

	codes := []uint8{0x0e, 0x18, 0x22, 0x2c, 0x40, 0x54, 0x68, 0x7c}

	for _, v := range codes {
		if code >= v && code <= v+4 {
			substraction = v - 4
			break
		}
	}

	paletteLength = code - substraction

	debug := false
	if debug {
		out := fmt.Sprintf("%v\t0x%X\t", stripNumber, code)
		if method == METHOD_ONE {
			out += "   1"
		} else if method == METHOD_TWO {
			out += "   2"
		}
		if direction == VERTICAL {
			out += "\tVert"
		} else {
			out += "\tHoriz"
		}
		out += "\t" + string(code-substraction)
		if transparent == TRANSP {
			out += "\tYes"
		} else {
			out += "\tNo"
		}
		log.Println(out)
	}
	return
}
<|endoftext|>


<commit_before>// Package imgcat provides helpers for embedding images (gif, png, jpeg)
// in the terminal output as suppored by iTerm 2.9+
// and documented at https://www.iterm2.com/images.html
package imgcat

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"image"
	"image/png"
	"io"
	"os"
)

// Cat embeds given io.Reader in the given io.Writer
func Cat(r io.Reader, w io.Writer) error {

	return embed(r, w)
}

// CatRGBA embeds given image.RGBA in the given io.Writer
func CatRGBA(i *image.RGBA, w io.Writer) error {

	b, err := imageAsPngBytes(i)
	if err != nil {
		return err
	}
	return embed(b, w)
}

// CatImage embeds given image.Image in the given io.Writer
func CatImage(i *image.Image, w io.Writer) error {

	b, err := imageAsPngBytes(*i)
	if err != nil {
		return err
	}
	return embed(b, w)
}

// CatFile embeds image file in the given io.Writer
func CatFile(fileName string, w io.Writer) error {

	r, err := os.Open(fileName)
	if err != nil {
		return err
	}

	return embed(r, w)
}

func embed(r io.Reader, w io.Writer) error {

	buf := new(bytes.Buffer)
	buf.ReadFrom(r)

	// tmux requires unrecognized OSC sequences to be wrapped with DCS tmux;
	// <sequence> ST, and for all ESCs in <sequence> to be replaced with ESC ESC. It
	// only accepts ESC backslash for ST.
	fmt.Fprint(w, "\033Ptmux;\033\033]")

	fmt.Fprintf(w, "1337;File=;inline=1:")

	encoder := base64.NewEncoder(base64.StdEncoding, w)
	defer encoder.Close()

	_, err := encoder.Write(buf.Bytes())
	if err != nil {
		return err
	}

	// More of the tmux workaround described above.
	fmt.Fprintln(w, "\a\033\\")

	return nil
}

func imageAsPngBytes(i image.Image) (io.Reader, error) {

	buf := new(bytes.Buffer)
	err := png.Encode(buf, i)
	if err != nil {
		return nil, err
	}
	return buf, nil
}
<commit_msg>simplify<commit_after>// Package imgcat provides helpers for embedding images (gif, png, jpeg)
// in the terminal output as suppored by iTerm 2.9+
// and documented at https://www.iterm2.com/images.html
package imgcat

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"image"
	"image/png"
	"io"
	"os"
)

// Cat embeds given io.Reader in the given io.Writer
func Cat(r io.Reader, w io.Writer) error {

	return embed(r, w)
}

// CatRGBA embeds given image.RGBA in the given io.Writer
func CatRGBA(i *image.RGBA, w io.Writer) error {

	b, err := imageAsPngBytes(i)
	if err != nil {
		return err
	}
	return embed(b, w)
}

// CatImage embeds given image.Image in the given io.Writer
func CatImage(i *image.Image, w io.Writer) error {

	b, err := imageAsPngBytes(*i)
	if err != nil {
		return err
	}
	return embed(b, w)
}

// CatFile embeds image file in the given io.Writer
func CatFile(fileName string, w io.Writer) error {

	r, err := os.Open(fileName)
	if err != nil {
		return err
	}

	return embed(r, w)
}

func embed(r io.Reader, w io.Writer) error {

	buf := new(bytes.Buffer)
	buf.ReadFrom(r)

	// tmux requires unrecognized OSC sequences to be wrapped with DCS tmux;
	// <sequence> ST, and for all ESCs in <sequence> to be replaced with ESC ESC. It
	// only accepts ESC backslash for ST.
	fmt.Fprint(w, "\033Ptmux;\033\033]1337;File=;inline=1:")

	encoder := base64.NewEncoder(base64.StdEncoding, w)
	defer encoder.Close()

	_, err := encoder.Write(buf.Bytes())
	if err != nil {
		return err
	}

	// More of the tmux workaround described above.
	fmt.Fprintln(w, "\a\033\\")
	return nil
}

func imageAsPngBytes(i image.Image) (io.Reader, error) {

	buf := new(bytes.Buffer)
	err := png.Encode(buf, i)
	if err != nil {
		return nil, err
	}
	return buf, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/koding/klient/cmd/klientctl/errors"
	"github.com/koding/klient/cmd/klientctl/util"
)

// MountCommand mounts a folder on remote machine to local folder by machine
// name.
func MountCommand(c *cli.Context) int {
	if len(c.Args()) < 2 {
		cli.ShowCommandHelp(c, "mount")
		return 1
	}

	var (
		name       = c.Args()[0]
		localPath  = c.Args()[1]
		remotePath = c.String("remotepath") // note the lowercase of all chars
	)

	// allow scp like declaration, ie `<machine name>:/path/to/remote`
	if strings.Contains(name, ":") {
		names := strings.Split(name, ":")
		name, remotePath = names[0], names[1]
	}

	// send absolute local path to klient unless local path is empty
	if strings.TrimSpace(localPath) != "" {
		absoluteLocalPath, err := filepath.Abs(localPath)
		if err == nil {
			localPath = absoluteLocalPath
		}
	}

	// Ask the user if they want the localPath created, if it does not exist.
	if err := askToCreate(localPath, os.Stdin, os.Stdout); err != nil {
		fmt.Printf(
			"Error: Unable to create specified localPath '%s'",
			localPath)
		return 1
	}

	// Check if the local path exists, and ask the user if they want to create it

	mountRequest := struct {
		Name       string `json:"name"`
		LocalPath  string `json:"localPath"`
		RemotePath string `json:"remotePath"`
	}{
		Name:      name,
		LocalPath: localPath,
	}

	// RemotePath is optional
	if remotePath != "" {
		mountRequest.RemotePath = remotePath
	}

	k, err := CreateKlientClient(NewKlientOptions())
	if err != nil {
		fmt.Printf("Error connecting to remove machine: '%s'\n", err)
		return 1
	}

	if err := k.Dial(); err != nil {
		fmt.Printf("Error connecting to remove machine: '%s'\n", err)
		return 1
	}

	resp, err := k.Tell("remote.mountFolder", mountRequest)
	if err != nil {
		fmt.Printf("Error mounting folder: '%s'\n", err)
		return 1
	}

	// response can be nil even when there's no err
	if resp != nil {
		var warning string
		if err := resp.Unmarshal(&warning); err != nil {
			return 0
		}

		if len(warning) > 0 {
			fmt.Printf("Warning: %s\n\n", warning)
		}
	}

	fmt.Println("Successfully mounted:", localPath)

	return 0
}

// askToCreate checks if the folder does not exist, and creates it
// if the user chooses to. If the user does *not* choose to create it,
// we return an IsNotExist error.
//
// TODO: Handle the case where a user types stuff in before being prompted,
// and then the prompt uses that. Ie, flush the input so that what we
// read is new input from the user. Not tested :)
func askToCreate(p string, r io.Reader, w io.Writer) error {
	_, err := os.Stat(p)

	// If we fail to stat the file, and it's *not* IsNotExist, we may be
	// having permission issues or some other related issue. Return
	// the error.
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	// If there was no error stating the path, it already exists -
	// we can return, as there's nothing we need to do.
	if err == nil {
		return nil
	}

	fmt.Fprintln(w,
		"The mount folder does not exist, would you like to create it? [Y/n]",
	)

	// To understand why we're creating a bReader here, please see
	// the docstring on YesNoConfirmWithDefault().
	bReader := bufio.NewReader(r)

	// Retry YesNo confirmation 3 times if needed
	var createFolder bool
	for i := 0; i < 3; i++ {
		createFolder, err = util.YesNoConfirmWithDefault(bReader, true)
		// If the user supplied an accepted value, stop trying
		if err == nil {
			break
		}
		// If err != nil, then the error did not provide an understood
		// response.
		fmt.Fprintln(w, "Invalid response, please type 'yes' or 'no'")
	}

	// If the retry loop exited with an error, the user failed to give
	// a meaningful response to the YesNo confirmation.
	if err != nil {
		return err
	}

	// The user chose not to create the folder. We cannot mount something that
	// doesn't exist - so we must fail here with an error.
	if !createFolder {
		return errs.ErrUserCancelled
	}

	err = os.Mkdir(p, 0655)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>styleguide: Return error<commit_after>package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/codegangsta/cli"
	"github.com/koding/klient/cmd/klientctl/errors"
	"github.com/koding/klient/cmd/klientctl/util"
)

// MountCommand mounts a folder on remote machine to local folder by machine
// name.
func MountCommand(c *cli.Context) int {
	if len(c.Args()) < 2 {
		cli.ShowCommandHelp(c, "mount")
		return 1
	}

	var (
		name       = c.Args()[0]
		localPath  = c.Args()[1]
		remotePath = c.String("remotepath") // note the lowercase of all chars
	)

	// allow scp like declaration, ie `<machine name>:/path/to/remote`
	if strings.Contains(name, ":") {
		names := strings.Split(name, ":")
		name, remotePath = names[0], names[1]
	}

	// send absolute local path to klient unless local path is empty
	if strings.TrimSpace(localPath) != "" {
		absoluteLocalPath, err := filepath.Abs(localPath)
		if err == nil {
			localPath = absoluteLocalPath
		}
	}

	// Ask the user if they want the localPath created, if it does not exist.
	if err := askToCreate(localPath, os.Stdin, os.Stdout); err != nil {
		fmt.Printf(
			"Error: Unable to create specified localPath '%s'",
			localPath)
		return 1
	}

	// Check if the local path exists, and ask the user if they want to create it

	mountRequest := struct {
		Name       string `json:"name"`
		LocalPath  string `json:"localPath"`
		RemotePath string `json:"remotePath"`
	}{
		Name:      name,
		LocalPath: localPath,
	}

	// RemotePath is optional
	if remotePath != "" {
		mountRequest.RemotePath = remotePath
	}

	k, err := CreateKlientClient(NewKlientOptions())
	if err != nil {
		fmt.Printf("Error connecting to remove machine: '%s'\n", err)
		return 1
	}

	if err := k.Dial(); err != nil {
		fmt.Printf("Error connecting to remove machine: '%s'\n", err)
		return 1
	}

	resp, err := k.Tell("remote.mountFolder", mountRequest)
	if err != nil {
		fmt.Printf("Error mounting folder: '%s'\n", err)
		return 1
	}

	// response can be nil even when there's no err
	if resp != nil {
		var warning string
		if err := resp.Unmarshal(&warning); err != nil {
			return 0
		}

		if len(warning) > 0 {
			fmt.Printf("Warning: %s\n\n", warning)
		}
	}

	fmt.Println("Successfully mounted:", localPath)

	return 0
}

// askToCreate checks if the folder does not exist, and creates it
// if the user chooses to. If the user does *not* choose to create it,
// we return an IsNotExist error.
//
// TODO: Handle the case where a user types stuff in before being prompted,
// and then the prompt uses that. Ie, flush the input so that what we
// read is new input from the user. Not tested :)
func askToCreate(p string, r io.Reader, w io.Writer) error {
	_, err := os.Stat(p)

	// If we fail to stat the file, and it's *not* IsNotExist, we may be
	// having permission issues or some other related issue. Return
	// the error.
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	// If there was no error stating the path, it already exists -
	// we can return, as there's nothing we need to do.
	if err == nil {
		return nil
	}

	fmt.Fprintln(w,
		"The mount folder does not exist, would you like to create it? [Y/n]",
	)

	// To understand why we're creating a bReader here, please see
	// the docstring on YesNoConfirmWithDefault().
	bReader := bufio.NewReader(r)

	// Retry YesNo confirmation 3 times if needed
	var createFolder bool
	for i := 0; i < 3; i++ {
		createFolder, err = util.YesNoConfirmWithDefault(bReader, true)
		// If the user supplied an accepted value, stop trying
		if err == nil {
			break
		}
		// If err != nil, then the error did not provide an understood
		// response.
		fmt.Fprintln(w, "Invalid response, please type 'yes' or 'no'")
	}

	// If the retry loop exited with an error, the user failed to give
	// a meaningful response to the YesNo confirmation.
	if err != nil {
		return err
	}

	// The user chose not to create the folder. We cannot mount something that
	// doesn't exist - so we must fail here with an error.
	if !createFolder {
		return errs.ErrUserCancelled
	}

	return os.Mkdir(p, 0655)
}
<|endoftext|>


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

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"time"

	"github.com/googlecloudplatform/gcsfuse/fs"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/fuse"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcscaching"
	"golang.org/x/net/context"
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s [flags] <mount-point>\n", os.Args[0])
	flag.PrintDefaults()
}

var fBucketName = flag.String("bucket", "", "Name of GCS bucket to mount.")

var fTempDir = flag.String(
	"temp_dir", "",
	"The temporary directory in which to store local copies of GCS objects. "+
		"If empty, the system default (probably /tmp) will be used.")

var fTempDirLimit = flag.Int64(
	"temp_dir_bytes", "",
	"A desired limit on the number of bytes used in --temp_dir. May be exceeded "+
		"for dirty files that have not been flushed or closed.")

var fImplicitDirs = flag.Bool(
	"implicit_dirs",
	false,
	"Implicitly define directories based on their content. See docs/semantics.md.")

var fSupportNlink = flag.Bool(
	"support_nlink",
	false,
	"Return meaningful values for nlink from fstat(2). See docs/semantics.md.")

var fStatCacheTTL = flag.String(
	"stat_cache_ttl",
	"",
	"If non-empty, a duration specifying how long to cache StatObject results "+
		"from GCS, e.g. \"2s\" or \"15ms\". See docs/semantics.md for more.")

var fTypeCacheTTL = flag.String(
	"type_cache_ttl",
	"",
	"If non-empty, a duration specifying how long to cache name -> file/dir type "+
		"mappings in directory inodes, e.g. \"2s\" or \"15ms\". "+
		"See docs/semantics.md.")

func getBucketName() string {
	s := *fBucketName
	if s == "" {
		fmt.Println("You must set --bucket.")
		os.Exit(1)
	}

	return s
}

func registerSIGINTHandler(mountPoint string) {
	// Register for SIGINT.
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)

	// Start a goroutine that will unmount when the signal is received.
	go func() {
		for {
			<-signalChan
			log.Println("Received SIGINT, attempting to unmount...")

			err := fuse.Unmount(mountPoint)
			if err != nil {
				log.Printf("Failed to unmount in response to SIGINT: %v", err)
			} else {
				log.Printf("Successfully unmounted in response to SIGINT.")
				return
			}
		}
	}()
}

func getBucket() (b gcs.Bucket) {
	// Set up a GCS connection.
	log.Println("Initializing GCS connection.")
	conn, err := getConn()
	if err != nil {
		log.Fatal("Couldn't get GCS connection: ", err)
	}

	// Extract the appropriate bucket.
	b = conn.GetBucket(getBucketName())

	// Enable cached StatObject results, if appropriate.
	if *fStatCacheTTL != "" {
		ttl, err := time.ParseDuration(*fStatCacheTTL)
		if err != nil {
			log.Fatalf("Invalid --stat_cache_ttl: %v", err)
			return
		}

		const cacheCapacity = 4096
		b = gcscaching.NewFastStatBucket(
			ttl,
			gcscaching.NewStatCache(cacheCapacity),
			timeutil.RealClock(),
			b)
	}

	return
}

func main() {
	// Make logging output better.
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)

	// Set up flags.
	flag.Usage = usage
	flag.Parse()

	// Grab the mount point.
	if flag.NArg() != 1 {
		usage()
		os.Exit(1)
	}

	mountPoint := flag.Arg(0)

	// Parse --type_cache_ttl
	var typeCacheTTL time.Duration
	if *fTypeCacheTTL != "" {
		var err error
		typeCacheTTL, err = time.ParseDuration(*fTypeCacheTTL)
		if err != nil {
			log.Fatalf("Invalid --type_cache_ttl: %v", err)
			return
		}
	}

	// Create a file system server.
	serverCfg := &fs.ServerConfig{
		Clock:               timeutil.RealClock(),
		Bucket:              getBucket(),
		TempDir:             *fTempDir,
		TempDirLimit:        *fTempDirLimit,
		ImplicitDirectories: *fImplicitDirs,
		SupportNlink:        *fSupportNlink,
		DirTypeCacheTTL:     typeCacheTTL,
	}

	server, err := fs.NewServer(serverCfg)
	if err != nil {
		log.Fatal("fs.NewServer:", err)
	}

	// Mount the file system.
	mountedFS, err := fuse.Mount(mountPoint, server, &fuse.MountConfig{})
	if err != nil {
		log.Fatal("Mount:", err)
	}

	log.Println("File system has been successfully mounted.")

	// Let the user unmount with Ctrl-C (SIGINT).
	registerSIGINTHandler(mountedFS.Dir())

	// Wait for it to be unmounted.
	if err := mountedFS.Join(context.Background()); err != nil {
		log.Fatal("MountedFileSystem.Join:", err)
	}

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

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"time"

	"github.com/googlecloudplatform/gcsfuse/fs"
	"github.com/googlecloudplatform/gcsfuse/timeutil"
	"github.com/jacobsa/fuse"
	"github.com/jacobsa/gcloud/gcs"
	"github.com/jacobsa/gcloud/gcs/gcscaching"
	"golang.org/x/net/context"
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "  %s [flags] <mount-point>\n", os.Args[0])
	flag.PrintDefaults()
}

var fBucketName = flag.String("bucket", "", "Name of GCS bucket to mount.")

var fTempDir = flag.String(
	"temp_dir", "",
	"The temporary directory in which to store local copies of GCS objects. "+
		"If empty, the system default (probably /tmp) will be used.")

var fTempDirLimit = flag.Int64(
	"temp_dir_bytes", 1<<31,
	"A desired limit on the number of bytes used in --temp_dir. May be exceeded "+
		"for dirty files that have not been flushed or closed.")

var fImplicitDirs = flag.Bool(
	"implicit_dirs",
	false,
	"Implicitly define directories based on their content. See docs/semantics.md.")

var fSupportNlink = flag.Bool(
	"support_nlink",
	false,
	"Return meaningful values for nlink from fstat(2). See docs/semantics.md.")

var fStatCacheTTL = flag.String(
	"stat_cache_ttl",
	"",
	"If non-empty, a duration specifying how long to cache StatObject results "+
		"from GCS, e.g. \"2s\" or \"15ms\". See docs/semantics.md for more.")

var fTypeCacheTTL = flag.String(
	"type_cache_ttl",
	"",
	"If non-empty, a duration specifying how long to cache name -> file/dir type "+
		"mappings in directory inodes, e.g. \"2s\" or \"15ms\". "+
		"See docs/semantics.md.")

func getBucketName() string {
	s := *fBucketName
	if s == "" {
		fmt.Println("You must set --bucket.")
		os.Exit(1)
	}

	return s
}

func registerSIGINTHandler(mountPoint string) {
	// Register for SIGINT.
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)

	// Start a goroutine that will unmount when the signal is received.
	go func() {
		for {
			<-signalChan
			log.Println("Received SIGINT, attempting to unmount...")

			err := fuse.Unmount(mountPoint)
			if err != nil {
				log.Printf("Failed to unmount in response to SIGINT: %v", err)
			} else {
				log.Printf("Successfully unmounted in response to SIGINT.")
				return
			}
		}
	}()
}

func getBucket() (b gcs.Bucket) {
	// Set up a GCS connection.
	log.Println("Initializing GCS connection.")
	conn, err := getConn()
	if err != nil {
		log.Fatal("Couldn't get GCS connection: ", err)
	}

	// Extract the appropriate bucket.
	b = conn.GetBucket(getBucketName())

	// Enable cached StatObject results, if appropriate.
	if *fStatCacheTTL != "" {
		ttl, err := time.ParseDuration(*fStatCacheTTL)
		if err != nil {
			log.Fatalf("Invalid --stat_cache_ttl: %v", err)
			return
		}

		const cacheCapacity = 4096
		b = gcscaching.NewFastStatBucket(
			ttl,
			gcscaching.NewStatCache(cacheCapacity),
			timeutil.RealClock(),
			b)
	}

	return
}

func main() {
	// Make logging output better.
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)

	// Set up flags.
	flag.Usage = usage
	flag.Parse()

	// Grab the mount point.
	if flag.NArg() != 1 {
		usage()
		os.Exit(1)
	}

	mountPoint := flag.Arg(0)

	// Parse --type_cache_ttl
	var typeCacheTTL time.Duration
	if *fTypeCacheTTL != "" {
		var err error
		typeCacheTTL, err = time.ParseDuration(*fTypeCacheTTL)
		if err != nil {
			log.Fatalf("Invalid --type_cache_ttl: %v", err)
			return
		}
	}

	// Create a file system server.
	serverCfg := &fs.ServerConfig{
		Clock:               timeutil.RealClock(),
		Bucket:              getBucket(),
		TempDir:             *fTempDir,
		TempDirLimit:        *fTempDirLimit,
		ImplicitDirectories: *fImplicitDirs,
		SupportNlink:        *fSupportNlink,
		DirTypeCacheTTL:     typeCacheTTL,
	}

	server, err := fs.NewServer(serverCfg)
	if err != nil {
		log.Fatal("fs.NewServer:", err)
	}

	// Mount the file system.
	mountedFS, err := fuse.Mount(mountPoint, server, &fuse.MountConfig{})
	if err != nil {
		log.Fatal("Mount:", err)
	}

	log.Println("File system has been successfully mounted.")

	// Let the user unmount with Ctrl-C (SIGINT).
	registerSIGINTHandler(mountedFS.Dir())

	// Wait for it to be unmounted.
	if err := mountedFS.Join(context.Background()); err != nil {
		log.Fatal("MountedFileSystem.Join:", err)
	}

	log.Println("Successfully exiting.")
}
<|endoftext|>


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

// Package mq provides interfaces for message-queue sockets.
package mq

import (
	"fmt"
	"strings"
	"sync"
)

// Socket is the main access handle that clients use to access the Fer system.
type Socket interface {
	// Close closes the open Socket
	Close() error

	// Send puts the message on the outbound send queue.
	// Send blocks until the message can be queued or the send deadline expires.
	Send(data []byte) error

	// Recv receives a complete message.
	Recv() ([]byte, error)

	// Listen connects alocal endpoint to the Socket.
	Listen(addr string) error

	// Dial connects a remote endpoint to the Socket.
	Dial(addr string) error

	// Type returns the type of this Socket (PUB, SUB, ...)
	Type() SocketType

	// GetOption is used to retrieve an option for a socket.
	//GetOption(name string) (interface{}, error)

	// SetOption is used to set an option for a socket.
	//SetOption(name string, value interface{}) error
}

// SocketType describes the type of a socket (PUB, SUB, PUSH, PULL, ...)
type SocketType int

// List of known socket types.
// Each Fer MQ driver ("zeromq", "nanomsg", ...) may support a different subset of
// socket types.
const (
	Invalid SocketType = iota
	Sub
	Pub
	XSub
	XPub
	Push
	Pull
	Req
	Rep
	Dealer
	Router
	Pair
	Bus
)

func (typ SocketType) String() string {
	switch typ {
	case Invalid:
		return "invalid"
	case Sub:
		return "sub"
	case Pub:
		return "pub"
	case XSub:
		return "xsub"
	case XPub:
		return "xpub"
	case Push:
		return "push"
	case Pull:
		return "pull"
	case Req:
		return "req"
	case Rep:
		return "rep"
	case Dealer:
		return "dealer"
	case Router:
		return "router"
	case Pair:
		return "pair"
	case Bus:
		return "bus"
	}
	return "N/A"
}

// SocketTypeFrom constructs a SocketType from its name.
// The matching is case insensitive.
// SocketTypeFrom panics if the given socket type name is invalid.
func SocketTypeFrom(name string) SocketType {
	switch strings.ToLower(name) {
	case "sub":
		return Sub
	case "pub":
		return Pub
	case "xpub":
		return XPub
	case "xsub":
		return XSub
	case "push":
		return Push
	case "pull":
		return Pull
	case "req":
		return Req
	case "rep":
		return Rep
	case "dealer":
		return Dealer
	case "router":
		return Router
	case "pair":
		return Pair
	case "bus":
		return Bus
	}
	panic(fmt.Errorf("fer: invalid socket type name (value=%q)", name))
}

var drivers struct {
	sync.RWMutex
	db map[string]Driver
}

// Register registers a new Fer MQ driver plugin
func Register(name string, drv Driver) {
	drivers.Lock()
	defer drivers.Unlock()
	if _, dup := drivers.db[name]; dup {
		panic(fmt.Errorf("fer: driver with name %q already registered", name))
	}
	drivers.db[name] = drv
}

// Open returns a previously registered driver plugin
//
// e.g.
//  zmq, err := fer.Open("zeromq")
//  nn,  err := fer.Open("nanomsg")
func Open(name string) (Driver, error) {
	drivers.RLock()
	defer drivers.RUnlock()
	drv, ok := drivers.db[name]
	if !ok {
		return nil, fmt.Errorf("fer: no such driver %q", name)
	}
	return drv, nil
}

// Driver is a Fer plugin to create FairMQ-compatible message queue communications
type Driver interface {
	NewSocket(typ SocketType) (Socket, error)
	Name() string
}

func init() {
	drivers.Lock()
	drivers.db = make(map[string]Driver)
	drivers.Unlock()
}
<commit_msg>mq: fix typo in Socket documentation<commit_after>// Copyright 2016 The fer Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package mq provides interfaces for message-queue sockets.
package mq

import (
	"fmt"
	"strings"
	"sync"
)

// Socket is the main access handle that clients use to access the Fer system.
type Socket interface {
	// Close closes the open Socket
	Close() error

	// Send puts the message on the outbound send queue.
	// Send blocks until the message can be queued or the send deadline expires.
	Send(data []byte) error

	// Recv receives a complete message.
	Recv() ([]byte, error)

	// Listen connects a local endpoint to the Socket.
	Listen(addr string) error

	// Dial connects a remote endpoint to the Socket.
	Dial(addr string) error

	// Type returns the type of this Socket (PUB, SUB, ...)
	Type() SocketType

	// GetOption is used to retrieve an option for a socket.
	//GetOption(name string) (interface{}, error)

	// SetOption is used to set an option for a socket.
	//SetOption(name string, value interface{}) error
}

// SocketType describes the type of a socket (PUB, SUB, PUSH, PULL, ...)
type SocketType int

// List of known socket types.
// Each Fer MQ driver ("zeromq", "nanomsg", ...) may support a different subset of
// socket types.
const (
	Invalid SocketType = iota
	Sub
	Pub
	XSub
	XPub
	Push
	Pull
	Req
	Rep
	Dealer
	Router
	Pair
	Bus
)

func (typ SocketType) String() string {
	switch typ {
	case Invalid:
		return "invalid"
	case Sub:
		return "sub"
	case Pub:
		return "pub"
	case XSub:
		return "xsub"
	case XPub:
		return "xpub"
	case Push:
		return "push"
	case Pull:
		return "pull"
	case Req:
		return "req"
	case Rep:
		return "rep"
	case Dealer:
		return "dealer"
	case Router:
		return "router"
	case Pair:
		return "pair"
	case Bus:
		return "bus"
	}
	return "N/A"
}

// SocketTypeFrom constructs a SocketType from its name.
// The matching is case insensitive.
// SocketTypeFrom panics if the given socket type name is invalid.
func SocketTypeFrom(name string) SocketType {
	switch strings.ToLower(name) {
	case "sub":
		return Sub
	case "pub":
		return Pub
	case "xpub":
		return XPub
	case "xsub":
		return XSub
	case "push":
		return Push
	case "pull":
		return Pull
	case "req":
		return Req
	case "rep":
		return Rep
	case "dealer":
		return Dealer
	case "router":
		return Router
	case "pair":
		return Pair
	case "bus":
		return Bus
	}
	panic(fmt.Errorf("fer: invalid socket type name (value=%q)", name))
}

var drivers struct {
	sync.RWMutex
	db map[string]Driver
}

// Register registers a new Fer MQ driver plugin
func Register(name string, drv Driver) {
	drivers.Lock()
	defer drivers.Unlock()
	if _, dup := drivers.db[name]; dup {
		panic(fmt.Errorf("fer: driver with name %q already registered", name))
	}
	drivers.db[name] = drv
}

// Open returns a previously registered driver plugin
//
// e.g.
//  zmq, err := fer.Open("zeromq")
//  nn,  err := fer.Open("nanomsg")
func Open(name string) (Driver, error) {
	drivers.RLock()
	defer drivers.RUnlock()
	drv, ok := drivers.db[name]
	if !ok {
		return nil, fmt.Errorf("fer: no such driver %q", name)
	}
	return drv, nil
}

// Driver is a Fer plugin to create FairMQ-compatible message queue communications
type Driver interface {
	NewSocket(typ SocketType) (Socket, error)
	Name() string
}

func init() {
	drivers.Lock()
	drivers.db = make(map[string]Driver)
	drivers.Unlock()
}
<|endoftext|>


<commit_before>package yiigo

import (
	"fmt"

	_ "github.com/go-sql-driver/mysql"
	"github.com/jmoiron/sqlx"
)

var dbmap map[string]*sqlx.DB

// initMySQL init db connections
func initMySQL() error {
	dbmap = make(map[string]*sqlx.DB)

	sections := childSections("mysql")

	for _, v := range sections {
		host := v.Key("host").MustString("localhost")
		port := v.Key("post").MustInt(3306)
		username := v.Key("username").MustString("root")
		password := v.Key("password").MustString("")
		database := v.Key("database").MustString("test")
		charset := v.Key("charset").MustString("utf8mb4")
		collection := v.Key("collection").MustString("utf8mb4_general_ci")

		dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&collation=%s&parseTime=True&loc=Local", username, password, host, port, database, charset, collection)

		db, err := sqlx.Open("mysql", dsn)

		if err != nil {
			db.Close()
			return err
		}

		db.SetMaxOpenConns(v.Key("maxOpenConns").MustInt(20))
		db.SetMaxIdleConns(v.Key("maxIdleConns").MustInt(10))

		err = db.Ping()

		if err != nil {
			db.Close()
			return err
		}

		dbmap[v.Name()] = db
	}

	return nil
}

// DB get a db connection
func DB(connection ...string) (*sqlx.DB, error) {
	conn := "default"

	if len(connection) > 0 {
		conn = connection[0]
	}

	db, ok := dbmap[fmt.Sprintf("mysql.%s", conn)]

	if !ok {
		return nil, fmt.Errorf("database %s is not connected", conn)
	}

	return db, nil
}
<commit_msg>update<commit_after>package yiigo

import (
	"errors"
	"fmt"
	"reflect"
	"strings"

	_ "github.com/go-sql-driver/mysql"
	"github.com/jmoiron/sqlx"
)

type MySQL struct {
	Connection string
	Table      string
}

var dbmap map[string]*sqlx.DB

// SQL expression
type expr struct {
	expr string
	args []interface{}
}

/**
 * 初始化DB
 */
func initMySQL() error {
	dbmap = make(map[string]*sqlx.DB)
	sections := getChildSections("mysql")

	for k, v := range sections {
		host := GetEnvString(v, "host", "localhost")
		port := GetEnvInt(v, "port", 3306)
		username := GetEnvString(v, "username", "root")
		password := GetEnvString(v, "password", "")
		database := GetEnvString(v, "database", "test")
		charset := GetEnvString(v, "charset", "utf8mb4")
		collection := GetEnvString(v, "collection", "utf8mb4_general_ci")

		dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&collation=%s&parseTime=True&loc=Local", username, password, host, port, database, charset, collection)
		db, err := sqlx.Open("mysql", dsn)

		if err != nil {
			db.Close()

			return err
		}

		db.SetMaxOpenConns(GetEnvInt(v, "maxOpenConns", 20))
		db.SetMaxIdleConns(GetEnvInt(v, "maxIdleConns", 10))

		err = db.Ping()

		if err != nil {
			db.Close()

			return err
		}

		dbmap[k] = db
	}

	return nil
}

/**
 * 获取db
 * @return *sqlx.DB
 */
func (m *MySQL) getDB() (*sqlx.DB, error) {
	connection := m.Connection

	if connection == "" {
		connection = "default"
	}

	db, ok := dbmap[connection]

	if !ok {
		return nil, fmt.Errorf("database %s is not connected", connection)
	}

	return db, nil
}

/**
 * 获取表前缀
 * @return string
 */
func (m *MySQL) getPrefix() string {
	connection := m.Connection

	if connection == "" {
		connection = "default"
	}

	prefix := GetEnvString(fmt.Sprintf("mysql.%s", connection), "prefix", "")

	return prefix
}

/**
 * Insert 插入
 * @param data X 插入数据
 * @return int64, error 新增记录ID
 */
func (m *MySQL) Insert(data X) (int64, error) {
	db, err := m.getDB()

	if err != nil {
		return 0, err
	}

	sql, binds := m.buildInsert(data)
	result, err := db.Exec(sql, binds...)

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	id, err := result.LastInsertId()

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	return id, nil
}

/**
 * BatchInsert 批量插入
 * @param columns []string 插入的字段
 * @param data []X 插入数据
 * @return int64, error 影响的行数
 */
func (m *MySQL) BatchInsert(columns []string, data []X) (int64, error) {
	db, err := m.getDB()

	if err != nil {
		return 0, err
	}

	sql, binds := m.buildBatchInsert(columns, data)
	result, err := db.Exec(sql, binds...)

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	rows, err := result.RowsAffected()

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	return rows, nil
}

/**
 * Update 更新
 * @param query X 查询条件
 * yiigo.X{
 *     where string WHERE条件语句
 *     binds []interface{} WHERE语句中 "?" 的绑定值
 * }
 * @param data X 更新字段
 * @return int64, error 影响的行数
 */
func (m *MySQL) Update(query X, data X) (int64, error) {
	db, err := m.getDB()

	if err != nil {
		return 0, err
	}

	sql, binds := m.buildUpdate(query, data)
	_sql, args, err := sqlx.In(sql, binds...)

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	result, err := db.Exec(_sql, args...)

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, _sql, args)
	}

	rows, err := result.RowsAffected()

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, _sql, args)
	}

	return rows, nil
}

/**
 * Count 获取记录数
 * @param query X 查询条件
 * yiigo.X{
 *     where string WHERE语句
 *     binds []interface{} WHERE语句中 "?" 的绑定值
 * }
 * @param columns ...string 聚合字段，默认为：*
 * @return error
 */
func (m *MySQL) Count(query X, columns ...string) (int, error) {
	db, err := m.getDB()

	if err != nil {
		return 0, err
	}

	if len(columns) > 0 {
		query["select"] = fmt.Sprintf("COUNT(%s)", columns[0])
	} else {
		query["select"] = "COUNT(*)"
	}

	sql, binds := m.buildQuery(query)
	_sql, args, err := sqlx.In(sql, binds...)

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	count := 0
	err = db.Get(&count, _sql, args...)

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	return count, nil
}

/**
 * FindOne 查询单条记录
 * @param query X 查询条件
 * yiigo.X{
 *     select string SELECT语句
 *     join string JOIN语句
 *     where string WHERE语句
 *     binds []interface{} WHERE语句中 "?" 的绑定值
 * }
 * @param data interface{} 查询数据 (struct指针)
 * @return error
 */
func (m *MySQL) FindOne(query X, data interface{}) error {
	db, err := m.getDB()

	if err != nil {
		return err
	}

	query["limit"] = 1

	sql, binds := m.buildQuery(query)
	_sql, args, err := sqlx.In(sql, binds...)

	if err != nil {
		return fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	err = db.Get(data, _sql, args...)

	if err != nil {
		if err.Error() == "sql: no rows in result set" {
			return errors.New("not found")
		}

		return fmt.Errorf("%v, SQL: %s, args: %v", err, _sql, args)
	}

	return nil
}

/**
 * Find 查询多条记录
 * @param query X 查询条件
 * yiigo.X{
 *     select string SELECT语句
 *     join []string JOIN语句
 *     where string WHERE语句
 *     group string GROUP BY语句
 *     order string ORDER BY语句
 *     limit int LIMIT语句
 *     offset int OFFSET语句
 *     binds []interface{} WHERE语句中 "?" 的绑定值
 * }
 * @param data interface{} 查询数据 (struct切片指针)
 * @return error
 */
func (m *MySQL) Find(query X, data interface{}) error {
	db, err := m.getDB()

	if err != nil {
		return err
	}

	sql, binds := m.buildQuery(query)
	_sql, args, err := sqlx.In(sql, binds...)

	if err != nil {
		return fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	err = db.Select(data, _sql, args...)

	if err != nil {
		return fmt.Errorf("%v, SQL: %s, args: %v", err, _sql, args)
	}

	return nil
}

/**
 * FindAll 查询所有记录
 * @param data interface{} 查询数据 (struct切片指针)
 * @param columns ...string 查询字段
 * @return error
 */
func (m *MySQL) FindAll(data interface{}, columns ...string) error {
	db, err := m.getDB()

	if err != nil {
		return err
	}

	query := X{}

	if len(columns) > 0 {
		query["select"] = strings.Join(columns, ",")
	}

	sql, binds := m.buildQuery(query)
	err = db.Select(data, sql, binds...)

	if err != nil {
		return fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	return nil
}

/**
 * FindBySQL SQL查询
 * @param sql string SQL查询语句
 * @parms binds []interface{} SQL绑定值
 * @param data interface{} 查询数据 (struct指针或struct切片指针)
 * @return error
 */
func (m *MySQL) FindBySQL(sql string, binds []interface{}, data interface{}) error {
	db, err := m.getDB()

	if err != nil {
		return err
	}

	_sql, args, err := sqlx.In(sql, binds...)

	if err != nil {
		return fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	v := reflect.Indirect(reflect.ValueOf(data))

	if v.Kind() == reflect.Slice {
		err = db.Select(data, _sql, args...)
	} else {
		err = db.Get(data, _sql, args...)
	}

	if err != nil {
		if err.Error() == "sql: no rows in result set" {
			return errors.New("not found")
		}

		return fmt.Errorf("%v, SQL: %s, args: %v", err, _sql, args)
	}

	return nil
}

/**
 * Delete 删除
 * @param query X 查询条件
 * yiigo.X{
 *     where string WHERE语句
 *     binds []interface{} WHERE语句中 "?" 的绑定值
 * }
 * @return int64, error 影响的行数
 */
func (m *MySQL) Delete(query X) (int64, error) {
	db, err := m.getDB()

	if err != nil {
		return 0, err
	}

	sql, binds := m.buildDelete(query)
	_sql, args, err := sqlx.In(sql, binds...)

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, sql, binds)
	}

	result, err := db.Exec(_sql, args...)

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, _sql, args)
	}

	rows, err := result.RowsAffected()

	if err != nil {
		return 0, fmt.Errorf("%v, SQL: %s, args: %v", err, _sql, args)
	}

	return rows, nil
}

/**
 * DoTransactions 事务处理
 * @param operations X 操作集合
 * [
 *     yiigo.X{
 *		   "type": "insert"
 *	 	   "table": string,
 *		   "data": yiigo.X,
 *     },
 *     yiigo.X{
 *		   "type": "batchInsert"
 *	 	   "table": string,
 *		   "columns": []string
 *		   "data": []yiigo.X,
 *     },
 *     yiigo.X{
 *	       "type": "update"
 *	       "query": yiigo.X{
 *	           "table": string,
 * 			   "where": string,
 *			   "binds": []interface{},
 *         },
 *		   "data": yiigo.X,
 *     },
 *	   yiigo.X{
 *		   "type": "delete"
 *	 	   "query": yiigo.X{
 *	 	   	   "table": string,
 * 			   "where": string,
 *			   "binds": []interface{},
 *         },
 *     },
 * ]
 * @return error
 */
func (m *MySQL) DoTransactions(operations []X) error {
	db, err := m.getDB()

	if err != nil {
		return err
	}

	tx, err := db.Begin()

	if err != nil {
		return err
	}

	errSQL := ""
	errArgs := []interface{}{}

	for _, opt := range operations {
		optType := ""

		if v, ok := opt["type"]; ok {
			optType = v.(string)
		} else {
			continue
		}

		switch optType {
		case "insert":
			table := []string{}
			data := X{}

			if v, ok := opt["table"]; ok {
				table = append(table, v.(string))
			}

			if v, ok := opt["data"]; ok {
				data = v.(X)
			}

			sql, binds := m.buildInsert(data, table...)
			_, err = tx.Exec(sql, binds...)

			if err != nil {
				errSQL = sql
				errArgs = binds

				break
			}
		case "batchInsert":
			table := []string{}
			columns := []string{}
			data := []X{}

			if v, ok := opt["table"]; ok {
				table = append(table, v.(string))
			}

			if v, ok := opt["columns"]; ok {
				columns = v.([]string)
			}

			if v, ok := opt["data"]; ok {
				data = v.([]X)
			}

			sql, binds := m.buildBatchInsert(columns, data, table...)
			_, err = tx.Exec(sql, binds...)

			if err != nil {
				errSQL = sql
				errArgs = binds

				break
			}
		case "update":
			query := X{}
			data := X{}

			if v, ok := opt["query"]; ok {
				query = v.(X)
			}

			if v, ok := opt["data"]; ok {
				data = v.(X)
			}

			sql, binds := m.buildUpdate(query, data)
			_sql, args, _ := sqlx.In(sql, binds...)
			_, err = tx.Exec(_sql, args...)

			if err != nil {
				errSQL = _sql
				errArgs = args

				break
			}
		case "delete":
			query := X{}

			if v, ok := opt["query"]; ok {
				query = v.(X)
			}

			sql, binds := m.buildDelete(query)
			_sql, args, _ := sqlx.In(sql, binds...)
			_, err = tx.Exec(_sql, args...)

			if err != nil {
				errSQL = _sql
				errArgs = args

				break
			}
		}

		if err != nil {
			break
		}
	}

	if err != nil {
		tx.Rollback()

		return fmt.Errorf("%v, SQL: %s, args: %v", err, errSQL, errArgs)
	}

	tx.Commit()

	return nil
}

/**
 * Expr SQL表达式
 * @param expression string 表达式，如：yiigo.Expr("price * ? + ?", 2, 100)
 * @param args ...interface{} 表达式中 "?" 的绑定值
 * @return *expr
 */
func Expr(expression string, args ...interface{}) *expr {
	return &expr{expr: expression, args: args}
}

/**
 * buildInsert 构建插入SQL
 * @param data X 插入数据
 * @param tables ...string 插入表
 * @return string, []interface{}
 */
func (m *MySQL) buildInsert(data X, tables ...string) (string, []interface{}) {
	if len(tables) == 0 {
		tables = append(tables, m.Table)
	}

	prefix := m.getPrefix()

	columns := []string{}
	placeholders := []string{}
	binds := []interface{}{}

	for k, v := range data {
		columns = append(columns, k)
		placeholders = append(placeholders, "?")
		binds = append(binds, v)
	}

	sql := fmt.Sprintf("INSERT INTO %s%s (%s) VALUES (%s)", prefix, tables[0], strings.Join(columns, ","), strings.Join(placeholders, ","))

	return sql, binds
}

/**
 * buildBatchInsert 构建批量插入SQL
 * @param columns []string 插入字段
 * @param data []X 插入数据
 * @param tables ...string 插入表
 * @return string, []interface{}
 */
func (m *MySQL) buildBatchInsert(columns []string, data []X, tables ...string) (string, []interface{}) {
	if len(tables) == 0 {
		tables = append(tables, m.Table)
	}

	prefix := m.getPrefix()

	placeholders := []string{}
	binds := []interface{}{}

	for _, v := range data {
		bindvars := []string{}

		for _, column := range columns {
			binds = append(binds, v[column])
			bindvars = append(bindvars, "?")
		}

		placeholders = append(placeholders, fmt.Sprintf("(%s)", strings.Join(bindvars, ",")))
	}

	sql := fmt.Sprintf("INSERT INTO %s%s (%s) VALUES %s", prefix, tables[0], strings.Join(columns, ","), strings.Join(placeholders, ","))

	return sql, binds
}

/**
 * buildUpdate 构建更新SQL
 * @param query X 查询条件
 * @param data X 更新数据
 * @return string, []interface{}
 */
func (m *MySQL) buildUpdate(query X, data X) (string, []interface{}) {
	table := m.Table
	prefix := m.getPrefix()

	clauses := []string{}
	sets := []string{}
	binds := []interface{}{}

	if v, ok := query["table"]; ok {
		table = v.(string)
	}

	clauses = append(clauses, fmt.Sprintf("UPDATE %s%s", prefix, table))

	for k, v := range data {
		if expr, ok := v.(*expr); ok {
			sets = append(sets, fmt.Sprintf("%s = %s", k, expr.expr))
			binds = append(binds, expr.args...)
		} else {
			sets = append(sets, fmt.Sprintf("%s = ?", k))
			binds = append(binds, v)
		}
	}

	clauses = append(clauses, fmt.Sprintf("SET %s", strings.Join(sets, ",")))

	if v, ok := query["where"]; ok {
		clauses = append(clauses, fmt.Sprintf("WHERE %s", v.(string)))
	}

	if v, ok := query["binds"]; ok {
		binds = append(binds, v.([]interface{})...)
	}

	sql := strings.Join(clauses, " ")

	return sql, binds
}

/**
 * buildQuery 构建查询SQL
 * @param query X 查询条件
 * @return string, []interface{}
 */
func (m *MySQL) buildQuery(query X) (string, []interface{}) {
	table := m.Table
	prefix := m.getPrefix()

	clauses := []string{}
	binds := []interface{}{}

	if v, ok := query["select"]; ok {
		clauses = append(clauses, fmt.Sprintf("SELECT %s", v.(string)))
	} else {
		clauses = append(clauses, "SELECT *")
	}

	if v, ok := query["table"]; ok {
		table = v.(string)
	}

	if v, ok := query["join"]; ok {
		clauses = append(clauses, fmt.Sprintf("FROM %s%s AS a", prefix, table))

		for _, join := range v.([]string) {
			clauses = append(clauses, join)
		}
	} else {
		clauses = append(clauses, fmt.Sprintf("FROM %s%s", prefix, table))
	}

	if v, ok := query["where"]; ok {
		clauses = append(clauses, fmt.Sprintf("WHERE %s", v.(string)))
	}

	if v, ok := query["group"]; ok {
		clauses = append(clauses, fmt.Sprintf("GROUP BY %s", v.(string)))
	}

	if v, ok := query["order"]; ok {
		clauses = append(clauses, fmt.Sprintf("ORDER BY %s", v.(string)))
	}

	if v, ok := query["limit"]; ok {
		clauses = append(clauses, fmt.Sprintf("LIMIT %d", v.(int)))
	}

	if v, ok := query["offset"]; ok {
		clauses = append(clauses, fmt.Sprintf("OFFSET %d", v.(int)))
	}

	if v, ok := query["binds"]; ok {
		binds = append(binds, v.([]interface{})...)
	}

	sql := strings.Join(clauses, " ")

	return sql, binds
}

/**
 * buildDelete 构建删除SQL
 * @param query X 查询条件
 * @return string, []interface{}
 */
func (m *MySQL) buildDelete(query X) (string, []interface{}) {
	table := m.Table
	prefix := m.getPrefix()

	clauses := []string{}
	binds := []interface{}{}

	if v, ok := query["table"]; ok {
		table = v.(string)
	}

	clauses = append(clauses, fmt.Sprintf("DELETE FROM %s%s", prefix, table))

	if v, ok := query["where"]; ok {
		clauses = append(clauses, fmt.Sprintf("WHERE %s", v.(string)))
	}

	if v, ok := query["binds"]; ok {
		binds = append(binds, v.([]interface{})...)
	}

	sql := strings.Join(clauses, " ")

	return sql, binds
}
<|endoftext|>


<commit_before>package appdb

import (
	"github.com/lib/pq"
	"golang.org/x/net/context"

	"chain/database/pg"
	"chain/errors"
	"chain/log"
)

// Bucket represents an indexed namespace inside of a wallet
type Bucket struct {
	ID    string   `json:"bucket_id"`
	Label string   `json:"label"`
	Index []uint32 `json:"bucket_index"`
}

// CreateBucket inserts a bucket database record
// for the given wallet,
// and returns the new Bucket.
func CreateBucket(ctx context.Context, walletID, label string) (*Bucket, error) {
	if label == "" {
		return nil, ErrBadLabel
	}

	bucket := &Bucket{Label: label}

	const attempts = 3
	for i := 0; i < attempts; i++ {
		const q = `
			WITH incr AS (
				UPDATE wallets
				SET
					buckets_count=buckets_count+1,
					next_bucket_index=next_bucket_index+1
				WHERE id=$1
				RETURNING (next_bucket_index - 1)
			)
			INSERT INTO buckets (wallet_id, key_index, label)
			VALUES ($1, (TABLE incr), $2)
			RETURNING id, key_index(key_index)
		`
		err := pg.FromContext(ctx).QueryRow(q, walletID, label).Scan(
			&bucket.ID,
			(*pg.Uint32s)(&bucket.Index),
		)
		if pqErr, ok := err.(*pq.Error); ok && pqErr.Code.Name() == "unique_violation" {
			// There was an (expected) unique index conflict.
			// It is safe to try again.
			// This happens when there is contention incrementing
			// the bucket index.
			log.Write(ctx, "attempt", i, "error", err)
			if i == attempts-1 {
				return nil, err
			}
			continue
		} else if err != nil {
			return nil, err
		}
		break
	}

	return bucket, nil
}

// BucketBalance fetches the balances of assets contained in this bucket.
// It returns a slice of Balances, where each Balance contains an asset ID,
// a confirmed balance, and a total balance. The total and confirmed balances
// are currently the same.
func BucketBalance(ctx context.Context, bucketID string) ([]*Balance, error) {
	q := `
		SELECT asset_id, sum(amount)::bigint
		FROM utxos
		WHERE bucket_id=$1
		GROUP BY asset_id
		ORDER BY asset_id
	`
	rows, err := pg.FromContext(ctx).Query(q, bucketID)
	if err != nil {
		return nil, errors.Wrap(err, "balance query")
	}
	defer rows.Close()
	var bals []*Balance

	for rows.Next() {
		var (
			assetID string
			bal     int64
		)
		err = rows.Scan(&assetID, &bal)
		if err != nil {
			return nil, errors.Wrap(err, "row scan")
		}
		b := &Balance{
			AssetID:   assetID,
			Total:     bal,
			Confirmed: bal,
		}
		bals = append(bals, b)
	}
	if err = rows.Err(); err != nil {
		return nil, errors.Wrap(err, "rows error")
	}
	return bals, err
}

// ListBuckets returns a list of buckets contained in the given wallet.
func ListBuckets(ctx context.Context, walletID string) ([]*Bucket, error) {
	q := `
		SELECT id, label, key_index(key_index)
		FROM buckets
		WHERE wallet_id = $1
		ORDER BY created_at
	`
	rows, err := pg.FromContext(ctx).Query(q, walletID)
	if err != nil {
		return nil, errors.Wrap(err, "select query")
	}
	defer rows.Close()

	var buckets []*Bucket
	for rows.Next() {
		b := new(Bucket)
		err = rows.Scan(
			&b.ID,
			&b.Label,
			(*pg.Uint32s)(&b.Index),
		)
		if err != nil {
			return nil, errors.Wrap(err, "row scan")
		}
		buckets = append(buckets, b)
	}

	if err = rows.Err(); err != nil {
		return nil, errors.Wrap(err, "end row scan loop")
	}

	return buckets, err
}
<commit_msg>appdb: bucket ID field should serialize as "id"<commit_after>package appdb

import (
	"github.com/lib/pq"
	"golang.org/x/net/context"

	"chain/database/pg"
	"chain/errors"
	"chain/log"
)

// Bucket represents an indexed namespace inside of a wallet
type Bucket struct {
	ID    string   `json:"id"`
	Label string   `json:"label"`
	Index []uint32 `json:"bucket_index"`
}

// CreateBucket inserts a bucket database record
// for the given wallet,
// and returns the new Bucket.
func CreateBucket(ctx context.Context, walletID, label string) (*Bucket, error) {
	if label == "" {
		return nil, ErrBadLabel
	}

	bucket := &Bucket{Label: label}

	const attempts = 3
	for i := 0; i < attempts; i++ {
		const q = `
			WITH incr AS (
				UPDATE wallets
				SET
					buckets_count=buckets_count+1,
					next_bucket_index=next_bucket_index+1
				WHERE id=$1
				RETURNING (next_bucket_index - 1)
			)
			INSERT INTO buckets (wallet_id, key_index, label)
			VALUES ($1, (TABLE incr), $2)
			RETURNING id, key_index(key_index)
		`
		err := pg.FromContext(ctx).QueryRow(q, walletID, label).Scan(
			&bucket.ID,
			(*pg.Uint32s)(&bucket.Index),
		)
		if pqErr, ok := err.(*pq.Error); ok && pqErr.Code.Name() == "unique_violation" {
			// There was an (expected) unique index conflict.
			// It is safe to try again.
			// This happens when there is contention incrementing
			// the bucket index.
			log.Write(ctx, "attempt", i, "error", err)
			if i == attempts-1 {
				return nil, err
			}
			continue
		} else if err != nil {
			return nil, err
		}
		break
	}

	return bucket, nil
}

// BucketBalance fetches the balances of assets contained in this bucket.
// It returns a slice of Balances, where each Balance contains an asset ID,
// a confirmed balance, and a total balance. The total and confirmed balances
// are currently the same.
func BucketBalance(ctx context.Context, bucketID string) ([]*Balance, error) {
	q := `
		SELECT asset_id, sum(amount)::bigint
		FROM utxos
		WHERE bucket_id=$1
		GROUP BY asset_id
		ORDER BY asset_id
	`
	rows, err := pg.FromContext(ctx).Query(q, bucketID)
	if err != nil {
		return nil, errors.Wrap(err, "balance query")
	}
	defer rows.Close()
	var bals []*Balance

	for rows.Next() {
		var (
			assetID string
			bal     int64
		)
		err = rows.Scan(&assetID, &bal)
		if err != nil {
			return nil, errors.Wrap(err, "row scan")
		}
		b := &Balance{
			AssetID:   assetID,
			Total:     bal,
			Confirmed: bal,
		}
		bals = append(bals, b)
	}
	if err = rows.Err(); err != nil {
		return nil, errors.Wrap(err, "rows error")
	}
	return bals, err
}

// ListBuckets returns a list of buckets contained in the given wallet.
func ListBuckets(ctx context.Context, walletID string) ([]*Bucket, error) {
	q := `
		SELECT id, label, key_index(key_index)
		FROM buckets
		WHERE wallet_id = $1
		ORDER BY created_at
	`
	rows, err := pg.FromContext(ctx).Query(q, walletID)
	if err != nil {
		return nil, errors.Wrap(err, "select query")
	}
	defer rows.Close()

	var buckets []*Bucket
	for rows.Next() {
		b := new(Bucket)
		err = rows.Scan(
			&b.ID,
			&b.Label,
			(*pg.Uint32s)(&b.Index),
		)
		if err != nil {
			return nil, errors.Wrap(err, "row scan")
		}
		buckets = append(buckets, b)
	}

	if err = rows.Err(); err != nil {
		return nil, errors.Wrap(err, "end row scan loop")
	}

	return buckets, err
}
<|endoftext|>


<commit_before>package sqlite3_test

import (
	"database/sql"
	"fmt"
	"math/rand"
	"regexp"
	"strconv"
	"sync"
	"testing"
	"time"
)

type Dialect int

const (
	SQLITE Dialect = iota
	POSTGRESQL
	MYSQL
)

type DB struct {
	*testing.T
	*sql.DB
	dialect Dialect
	once    sync.Once
}

var db *DB

// the following tables will be created and dropped during the test
var testTables = []string{"foo", "bar", "t", "bench"}

var tests = []testing.InternalTest{
	{"TestBlobs", TestBlobs},
	{"TestManyQueryRow", TestManyQueryRow},
	{"TestTxQuery", TestTxQuery},
	{"TestPreparedStmt", TestPreparedStmt},
}

var benchmarks = []testing.InternalBenchmark{
	{"BenchmarkExec", BenchmarkExec},
	{"BenchmarkQuery", BenchmarkQuery},
	{"BenchmarkParams", BenchmarkParams},
	{"BenchmarkStmt", BenchmarkStmt},
	{"BenchmarkRows", BenchmarkRows},
	{"BenchmarkStmtRows", BenchmarkStmtRows},
}

// RunTests runs the SQL test suite
func RunTests(t *testing.T, d *sql.DB, dialect Dialect) {
	db = &DB{t, d, dialect, sync.Once{}}
	testing.RunTests(func(string, string) (bool, error) { return true, nil }, tests)

	if !testing.Short() {
		for _, b := range benchmarks {
			fmt.Printf("%-20s", b.Name)
			r := testing.Benchmark(b.F)
			fmt.Printf("%10d %10.0f req/s\n", r.N, float64(r.N)/r.T.Seconds())
		}
	}
	db.tearDown()
}

func (db *DB) mustExec(sql string, args ...interface{}) sql.Result {
	res, err := db.Exec(sql, args...)
	if err != nil {
		db.Fatalf("Error running %q: %v", sql, err)
	}
	return res
}

func (db *DB) tearDown() {
	for _, tbl := range testTables {
		switch db.dialect {
		case SQLITE:
			db.mustExec("drop table if exists " + tbl)
		case MYSQL, POSTGRESQL:
			db.mustExec("drop table if exists " + tbl)
		default:
			db.Fatal("unkown dialect")
		}
	}
}

// q replaces ? parameters if needed
func (db *DB) q(sql string) string {
	switch db.dialect {
	case POSTGRESQL: // repace with $1, $2, ..
		qrx := regexp.MustCompile(`\?`)
		n := 0
		return qrx.ReplaceAllStringFunc(sql, func(string) string {
			n++
			return "$" + strconv.Itoa(n)
		})
	}
	return sql
}

func (db *DB) blobType(size int) string {
	switch db.dialect {
	case SQLITE:
		return fmt.Sprintf("blob[%d]", size)
	case POSTGRESQL:
		return "bytea"
	case MYSQL:
		return fmt.Sprintf("VARBINARY(%d)", size)
	}
	panic("unkown dialect")
}

func (db *DB) serialPK() string {
	switch db.dialect {
	case SQLITE:
		return "integer primary key autoincrement"
	case POSTGRESQL:
		return "serial primary key"
	case MYSQL:
		return "integer primary key auto_increment"
	}
	panic("unkown dialect")
}

func (db *DB) now() string {
	switch db.dialect {
	case SQLITE:
		return "datetime('now')"
	case POSTGRESQL:
		return "now()"
	case MYSQL:
		return "now()"
	}
	panic("unkown dialect")
}

func makeBench() {
	if _, err := db.Exec("create table bench (n varchar(32), i integer, d double, s varchar(32), t datetime)"); err != nil {
		panic(err)
	}
	st, err := db.Prepare("insert into bench values (?, ?, ?, ?, ?)")
	if err != nil {
		panic(err)
	}
	defer st.Close()
	for i := 0; i < 100; i++ {
		if _, err = st.Exec(nil, i, float64(i), fmt.Sprintf("%d", i), time.Now()); err != nil {
			panic(err)
		}
	}
}

func TestResult(t *testing.T) {
	db.tearDown()
	db.mustExec("create temporary table test (id " + db.serialPK() + ", name varchar(10))")

	for i := 1; i < 3; i++ {
		r := db.mustExec(db.q("insert into test (name) values (?)"), fmt.Sprintf("row %d", i))
		n, err := r.RowsAffected()
		if err != nil {
			t.Fatal(err)
		}
		if n != 1 {
			t.Errorf("got %v, want %v", n, 1)
		}
		n, err = r.LastInsertId()
		if err != nil {
			t.Fatal(err)
		}
		if n != int64(i) {
			t.Errorf("got %v, want %v", n, i)
		}
	}
	if _, err := db.Exec("error!"); err == nil {
		t.Fatalf("expected error")
	}
}

func TestBlobs(t *testing.T) {
	db.tearDown()
	var blob = []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
	db.mustExec("create table foo (id integer primary key, bar " + db.blobType(16) + ")")
	db.mustExec(db.q("insert into foo (id, bar) values(?,?)"), 0, blob)

	want := fmt.Sprintf("%x", blob)

	b := make([]byte, 16)
	err := db.QueryRow(db.q("select bar from foo where id = ?"), 0).Scan(&b)
	got := fmt.Sprintf("%x", b)
	if err != nil {
		t.Errorf("[]byte scan: %v", err)
	} else if got != want {
		t.Errorf("for []byte, got %q; want %q", got, want)
	}

	err = db.QueryRow(db.q("select bar from foo where id = ?"), 0).Scan(&got)
	want = string(blob)
	if err != nil {
		t.Errorf("string scan: %v", err)
	} else if got != want {
		t.Errorf("for string, got %q; want %q", got, want)
	}
}

func TestManyQueryRow(t *testing.T) {
	if testing.Short() {
		t.Log("skipping in short mode")
		return
	}
	db.tearDown()
	db.mustExec("create table foo (id integer primary key, name varchar(50))")
	db.mustExec(db.q("insert into foo (id, name) values(?,?)"), 1, "bob")
	var name string
	for i := 0; i < 10000; i++ {
		err := db.QueryRow(db.q("select name from foo where id = ?"), 1).Scan(&name)
		if err != nil || name != "bob" {
			t.Fatalf("on query %d: err=%v, name=%q", i, err, name)
		}
	}
}

func TestTxQuery(t *testing.T) {
	db.tearDown()
	tx, err := db.Begin()
	if err != nil {
		t.Fatal(err)
	}
	defer tx.Rollback()

	_, err = tx.Exec("create table foo (id integer primary key, name varchar(50))")
	if err != nil {
		t.Fatal(err)
	}

	_, err = tx.Exec(db.q("insert into foo (id, name) values(?,?)"), 1, "bob")
	if err != nil {
		t.Fatal(err)
	}

	r, err := tx.Query(db.q("select name from foo where id = ?"), 1)
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()

	if !r.Next() {
		if r.Err() != nil {
			t.Fatal(err)
		}
		t.Fatal("expected one rows")
	}

	var name string
	err = r.Scan(&name)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPreparedStmt(t *testing.T) {
	db.tearDown()
	db.mustExec("CREATE TABLE t (count INT)")
	sel, err := db.Prepare("SELECT count FROM t ORDER BY count DESC")
	if err != nil {
		t.Fatalf("prepare 1: %v", err)
	}
	ins, err := db.Prepare(db.q("INSERT INTO t (count) VALUES (?)"))
	if err != nil {
		t.Fatalf("prepare 2: %v", err)
	}

	for n := 1; n <= 3; n++ {
		if _, err := ins.Exec(n); err != nil {
			t.Fatalf("insert(%d) = %v", n, err)
		}
	}

	const nRuns = 10
	ch := make(chan bool)
	for i := 0; i < nRuns; i++ {
		go func() {
			defer func() {
				ch <- true
			}()
			for j := 0; j < 10; j++ {
				count := 0
				if err := sel.QueryRow().Scan(&count); err != nil && err != sql.ErrNoRows {
					t.Errorf("Query: %v", err)
					return
				}
				if _, err := ins.Exec(rand.Intn(100)); err != nil {
					t.Errorf("Insert: %v", err)
					return
				}
			}
		}()
	}
	for i := 0; i < nRuns; i++ {
		<-ch
	}
}

// Benchmarks need to use panic() since b.Error errors are lost when
// running via testing.Benchmark() I would like to run these via go
// test -bench but calling Benchmark() from a benchmark test
// currently hangs go.

func BenchmarkExec(b *testing.B) {
	for i := 0; i < b.N; i++ {
		if _, err := db.Exec("select 1"); err != nil {
			panic(err)
		}
	}
}

func BenchmarkQuery(b *testing.B) {
	for i := 0; i < b.N; i++ {
		var n sql.NullString
		var i int
		var f float64
		var s string
//		var t time.Time
		if err := db.QueryRow("select null, 1, 1.1, 'foo'").Scan(&n, &i, &f, &s); err != nil {
			panic(err)
		}
	}
}

func BenchmarkParams(b *testing.B) {
	for i := 0; i < b.N; i++ {
		var n sql.NullString
		var i int
		var f float64
		var s string
//		var t time.Time
		if err := db.QueryRow("select ?, ?, ?, ?", nil, 1, 1.1, "foo").Scan(&n, &i, &f, &s); err != nil {
			panic(err)
		}
	}
}

func BenchmarkStmt(b *testing.B) {
	st, err := db.Prepare("select ?, ?, ?, ?")
	if err != nil {
		panic(err)
	}
	defer st.Close()

	for n := 0; n < b.N; n++ {
		var n sql.NullString
		var i int
		var f float64
		var s string
//		var t time.Time
		if err := st.QueryRow(nil, 1, 1.1, "foo").Scan(&n, &i, &f, &s); err != nil {
			panic(err)
		}
	}
}

func BenchmarkRows(b *testing.B) {
	db.once.Do(makeBench)

	for n := 0; n < b.N; n++ {
		var n sql.NullString
		var i int
		var f float64
		var s string
		var t time.Time
		r, err := db.Query("select * from bench")
		if err != nil {
			panic(err)
		}
		for r.Next() {
			if err = r.Scan(&n, &i, &f, &s, &t); err != nil {
				panic(err)
			}
		}
		if err = r.Err(); err != nil {
			panic(err)
		}
	}
}

func BenchmarkStmtRows(b *testing.B) {
	db.once.Do(makeBench)

	st, err := db.Prepare("select * from bench")
	if err != nil {
		panic(err)
	}
	defer st.Close()

	for n := 0; n < b.N; n++ {
		var n sql.NullString
		var i int
		var f float64
		var s string
		var t time.Time
		r, err := st.Query()
		if err != nil {
			panic(err)
		}
		for r.Next() {
			if err = r.Scan(&n, &i, &f, &s, &t); err != nil {
				panic(err)
			}
		}
		if err = r.Err(); err != nil {
			panic(err)
		}
	}
}
<commit_msg>Cosmetics: use WaitGroup instead of channel in test<commit_after>package sqlite3_test

import (
	"database/sql"
	"fmt"
	"math/rand"
	"regexp"
	"strconv"
	"sync"
	"testing"
	"time"
)

type Dialect int

const (
	SQLITE Dialect = iota
	POSTGRESQL
	MYSQL
)

type DB struct {
	*testing.T
	*sql.DB
	dialect Dialect
	once    sync.Once
}

var db *DB

// the following tables will be created and dropped during the test
var testTables = []string{"foo", "bar", "t", "bench"}

var tests = []testing.InternalTest{
	{"TestBlobs", TestBlobs},
	{"TestManyQueryRow", TestManyQueryRow},
	{"TestTxQuery", TestTxQuery},
	{"TestPreparedStmt", TestPreparedStmt},
}

var benchmarks = []testing.InternalBenchmark{
	{"BenchmarkExec", BenchmarkExec},
	{"BenchmarkQuery", BenchmarkQuery},
	{"BenchmarkParams", BenchmarkParams},
	{"BenchmarkStmt", BenchmarkStmt},
	{"BenchmarkRows", BenchmarkRows},
	{"BenchmarkStmtRows", BenchmarkStmtRows},
}

// RunTests runs the SQL test suite
func RunTests(t *testing.T, d *sql.DB, dialect Dialect) {
	db = &DB{t, d, dialect, sync.Once{}}
	testing.RunTests(func(string, string) (bool, error) { return true, nil }, tests)

	if !testing.Short() {
		for _, b := range benchmarks {
			fmt.Printf("%-20s", b.Name)
			r := testing.Benchmark(b.F)
			fmt.Printf("%10d %10.0f req/s\n", r.N, float64(r.N)/r.T.Seconds())
		}
	}
	db.tearDown()
}

func (db *DB) mustExec(sql string, args ...interface{}) sql.Result {
	res, err := db.Exec(sql, args...)
	if err != nil {
		db.Fatalf("Error running %q: %v", sql, err)
	}
	return res
}

func (db *DB) tearDown() {
	for _, tbl := range testTables {
		switch db.dialect {
		case SQLITE:
			db.mustExec("drop table if exists " + tbl)
		case MYSQL, POSTGRESQL:
			db.mustExec("drop table if exists " + tbl)
		default:
			db.Fatal("unkown dialect")
		}
	}
}

// q replaces ? parameters if needed
func (db *DB) q(sql string) string {
	switch db.dialect {
	case POSTGRESQL: // repace with $1, $2, ..
		qrx := regexp.MustCompile(`\?`)
		n := 0
		return qrx.ReplaceAllStringFunc(sql, func(string) string {
			n++
			return "$" + strconv.Itoa(n)
		})
	}
	return sql
}

func (db *DB) blobType(size int) string {
	switch db.dialect {
	case SQLITE:
		return fmt.Sprintf("blob[%d]", size)
	case POSTGRESQL:
		return "bytea"
	case MYSQL:
		return fmt.Sprintf("VARBINARY(%d)", size)
	}
	panic("unkown dialect")
}

func (db *DB) serialPK() string {
	switch db.dialect {
	case SQLITE:
		return "integer primary key autoincrement"
	case POSTGRESQL:
		return "serial primary key"
	case MYSQL:
		return "integer primary key auto_increment"
	}
	panic("unkown dialect")
}

func (db *DB) now() string {
	switch db.dialect {
	case SQLITE:
		return "datetime('now')"
	case POSTGRESQL:
		return "now()"
	case MYSQL:
		return "now()"
	}
	panic("unkown dialect")
}

func makeBench() {
	if _, err := db.Exec("create table bench (n varchar(32), i integer, d double, s varchar(32), t datetime)"); err != nil {
		panic(err)
	}
	st, err := db.Prepare("insert into bench values (?, ?, ?, ?, ?)")
	if err != nil {
		panic(err)
	}
	defer st.Close()
	for i := 0; i < 100; i++ {
		if _, err = st.Exec(nil, i, float64(i), fmt.Sprintf("%d", i), time.Now()); err != nil {
			panic(err)
		}
	}
}

func TestResult(t *testing.T) {
	db.tearDown()
	db.mustExec("create temporary table test (id " + db.serialPK() + ", name varchar(10))")

	for i := 1; i < 3; i++ {
		r := db.mustExec(db.q("insert into test (name) values (?)"), fmt.Sprintf("row %d", i))
		n, err := r.RowsAffected()
		if err != nil {
			t.Fatal(err)
		}
		if n != 1 {
			t.Errorf("got %v, want %v", n, 1)
		}
		n, err = r.LastInsertId()
		if err != nil {
			t.Fatal(err)
		}
		if n != int64(i) {
			t.Errorf("got %v, want %v", n, i)
		}
	}
	if _, err := db.Exec("error!"); err == nil {
		t.Fatalf("expected error")
	}
}

func TestBlobs(t *testing.T) {
	db.tearDown()
	var blob = []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}
	db.mustExec("create table foo (id integer primary key, bar " + db.blobType(16) + ")")
	db.mustExec(db.q("insert into foo (id, bar) values(?,?)"), 0, blob)

	want := fmt.Sprintf("%x", blob)

	b := make([]byte, 16)
	err := db.QueryRow(db.q("select bar from foo where id = ?"), 0).Scan(&b)
	got := fmt.Sprintf("%x", b)
	if err != nil {
		t.Errorf("[]byte scan: %v", err)
	} else if got != want {
		t.Errorf("for []byte, got %q; want %q", got, want)
	}

	err = db.QueryRow(db.q("select bar from foo where id = ?"), 0).Scan(&got)
	want = string(blob)
	if err != nil {
		t.Errorf("string scan: %v", err)
	} else if got != want {
		t.Errorf("for string, got %q; want %q", got, want)
	}
}

func TestManyQueryRow(t *testing.T) {
	if testing.Short() {
		t.Log("skipping in short mode")
		return
	}
	db.tearDown()
	db.mustExec("create table foo (id integer primary key, name varchar(50))")
	db.mustExec(db.q("insert into foo (id, name) values(?,?)"), 1, "bob")
	var name string
	for i := 0; i < 10000; i++ {
		err := db.QueryRow(db.q("select name from foo where id = ?"), 1).Scan(&name)
		if err != nil || name != "bob" {
			t.Fatalf("on query %d: err=%v, name=%q", i, err, name)
		}
	}
}

func TestTxQuery(t *testing.T) {
	db.tearDown()
	tx, err := db.Begin()
	if err != nil {
		t.Fatal(err)
	}
	defer tx.Rollback()

	_, err = tx.Exec("create table foo (id integer primary key, name varchar(50))")
	if err != nil {
		t.Fatal(err)
	}

	_, err = tx.Exec(db.q("insert into foo (id, name) values(?,?)"), 1, "bob")
	if err != nil {
		t.Fatal(err)
	}

	r, err := tx.Query(db.q("select name from foo where id = ?"), 1)
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()

	if !r.Next() {
		if r.Err() != nil {
			t.Fatal(err)
		}
		t.Fatal("expected one rows")
	}

	var name string
	err = r.Scan(&name)
	if err != nil {
		t.Fatal(err)
	}
}

func TestPreparedStmt(t *testing.T) {
	db.tearDown()
	db.mustExec("CREATE TABLE t (count INT)")
	sel, err := db.Prepare("SELECT count FROM t ORDER BY count DESC")
	if err != nil {
		t.Fatalf("prepare 1: %v", err)
	}
	ins, err := db.Prepare(db.q("INSERT INTO t (count) VALUES (?)"))
	if err != nil {
		t.Fatalf("prepare 2: %v", err)
	}

	for n := 1; n <= 3; n++ {
		if _, err := ins.Exec(n); err != nil {
			t.Fatalf("insert(%d) = %v", n, err)
		}
	}

	const nRuns = 10
	var wg sync.WaitGroup
	for i := 0; i < nRuns; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := 0; j < 10; j++ {
				count := 0
				if err := sel.QueryRow().Scan(&count); err != nil && err != sql.ErrNoRows {
					t.Errorf("Query: %v", err)
					return
				}
				if _, err := ins.Exec(rand.Intn(100)); err != nil {
					t.Errorf("Insert: %v", err)
					return
				}
			}
		}()
	}
	wg.Wait()
}

// Benchmarks need to use panic() since b.Error errors are lost when
// running via testing.Benchmark() I would like to run these via go
// test -bench but calling Benchmark() from a benchmark test
// currently hangs go.

func BenchmarkExec(b *testing.B) {
	for i := 0; i < b.N; i++ {
		if _, err := db.Exec("select 1"); err != nil {
			panic(err)
		}
	}
}

func BenchmarkQuery(b *testing.B) {
	for i := 0; i < b.N; i++ {
		var n sql.NullString
		var i int
		var f float64
		var s string
//		var t time.Time
		if err := db.QueryRow("select null, 1, 1.1, 'foo'").Scan(&n, &i, &f, &s); err != nil {
			panic(err)
		}
	}
}

func BenchmarkParams(b *testing.B) {
	for i := 0; i < b.N; i++ {
		var n sql.NullString
		var i int
		var f float64
		var s string
//		var t time.Time
		if err := db.QueryRow("select ?, ?, ?, ?", nil, 1, 1.1, "foo").Scan(&n, &i, &f, &s); err != nil {
			panic(err)
		}
	}
}

func BenchmarkStmt(b *testing.B) {
	st, err := db.Prepare("select ?, ?, ?, ?")
	if err != nil {
		panic(err)
	}
	defer st.Close()

	for n := 0; n < b.N; n++ {
		var n sql.NullString
		var i int
		var f float64
		var s string
//		var t time.Time
		if err := st.QueryRow(nil, 1, 1.1, "foo").Scan(&n, &i, &f, &s); err != nil {
			panic(err)
		}
	}
}

func BenchmarkRows(b *testing.B) {
	db.once.Do(makeBench)

	for n := 0; n < b.N; n++ {
		var n sql.NullString
		var i int
		var f float64
		var s string
		var t time.Time
		r, err := db.Query("select * from bench")
		if err != nil {
			panic(err)
		}
		for r.Next() {
			if err = r.Scan(&n, &i, &f, &s, &t); err != nil {
				panic(err)
			}
		}
		if err = r.Err(); err != nil {
			panic(err)
		}
	}
}

func BenchmarkStmtRows(b *testing.B) {
	db.once.Do(makeBench)

	st, err := db.Prepare("select * from bench")
	if err != nil {
		panic(err)
	}
	defer st.Close()

	for n := 0; n < b.N; n++ {
		var n sql.NullString
		var i int
		var f float64
		var s string
		var t time.Time
		r, err := st.Query()
		if err != nil {
			panic(err)
		}
		for r.Next() {
			if err = r.Scan(&n, &i, &f, &s, &t); err != nil {
				panic(err)
			}
		}
		if err = r.Err(); err != nil {
			panic(err)
		}
	}
}
<|endoftext|>