<commit_before>package internal

import (
	"go/ast"
	"log"
	"os"
	"path/filepath"
)

type Context struct {
	TranslatedassertImport *ast.Ident
	AssertImport           *ast.Ident
}

var (
	CacheDir   = filepath.Join(os.Getenv(homeEnv), ".gopwtcache")
	Testdata   = "testdata"
	TermWidth  = 0
	WorkingDir = ""
	Verbose    = false
)

func Translate(gopath string, importpath, fpath string) error {
	if debugLog {
		log.Printf("path=%s", fpath)
		log.Printf("gopath=%s", gopath)
		log.Printf("importpath=%s", importpath)
	}

	srcDir := filepath.Join(gopath, "src")
	err := os.MkdirAll(filepath.Join(srcDir, importpath), os.ModePerm)
	if err != nil {
		return err
	}

	return NewPackageContext(fpath, importpath, srcDir).Translate()
}
<commit_msg>fix(translator/internal): split cachedir by version<commit_after>package internal

import (
	"go/ast"
	"log"
	"os"
	"path/filepath"
	"runtime"
)

type Context struct {
	TranslatedassertImport *ast.Ident
	AssertImport           *ast.Ident
}

var (
	CacheDir   = filepath.Join(os.Getenv(homeEnv), ".gopwtcache", runtime.Version())
	Testdata   = "testdata"
	TermWidth  = 0
	WorkingDir = ""
	Verbose    = false
)

func Translate(gopath string, importpath, fpath string) error {
	if debugLog {
		log.Printf("path=%s", fpath)
		log.Printf("gopath=%s", gopath)
		log.Printf("importpath=%s", importpath)
	}

	srcDir := filepath.Join(gopath, "src")
	err := os.MkdirAll(filepath.Join(srcDir, importpath), os.ModePerm)
	if err != nil {
		return err
	}

	return NewPackageContext(fpath, importpath, srcDir).Translate()
}
<|endoftext|>


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

package skiplist_test

import (
	"fmt"
	"github.com/someonegg/gocontainer/skiplist"
	//"math/rand"
	//"time"
)

func elemCompare(l, r skiplist.Scorable) int {
	vl := l.(int)
	vr := r.(int)
	return vl - vr
}

func Example() {
	//rand.Seed(time.Now().Unix())

	// An empty skiplist and put some numbers in it.
	l := skiplist.NewList(elemCompare)

	// add and print rank
	e1 := l.Add(4)
	fmt.Println(l.Rank(e1))

	e2 := l.Add(1)
	fmt.Println(l.Rank(e2))

	e3 := l.Add(3)
	fmt.Println(l.Rank(e3))

	e4 := l.Add(2)
	fmt.Println(l.Rank(e4))

	e5 := l.Add(4)
	fmt.Println(l.Rank(e5))

	e6 := l.Add(6)
	fmt.Println(l.Rank(e6))

	fmt.Println()

	// print list after add.
	for e := l.Front(); e != nil; e = e.Next() {
		fmt.Println(e.Value)
	}

	fmt.Println()

	l.Remove(l.Find(6))
	l.Remove(l.Get(4))

	// print list after remove.
	for e := l.Front(); e != nil; e = e.Next() {
		fmt.Println(e.Value)
	}

	// Output:
	// 0
	// 0
	// 1
	// 1
	// 3
	// 5
	//
	// 1
	// 2
	// 3
	// 4
	// 4
	// 6
	//
	// 1
	// 2
	// 3
	// 4
}
<commit_msg>update test<commit_after>// Copyright 2015 someonegg. All rights reserscoreed.
// Use of this source code is goscoreerned by a BSD-style
// license that can be found in the LICENSE file.

package skiplist_test

import (
	"fmt"
	"github.com/someonegg/gocontainer/skiplist"
	//"math/rand"
	//"time"
)

type item struct {
	score int

	// other...
}

func itemCompare(l, r skiplist.Scorable) int {
	scoreL := int(0)
	switch v := l.(type) {
	case int:
		scoreL = v
	case *item:
		scoreL = v.score
	default:
		panic(l)
	}
	scoreR := int(0)
	switch v := r.(type) {
	case int:
		scoreR = v
	case *item:
		scoreR = v.score
	default:
		panic(r)
	}
	return scoreL - scoreR
}

func Example() {
	//rand.Seed(time.Now().Unix())

	// An empty skiplist and put some numbers in it.
	l := skiplist.NewList(itemCompare)

	// add and print rank
	e1 := l.Add(&item{score: 4})
	fmt.Println(l.Rank(e1))

	e2 := l.Add(&item{score: 1})
	fmt.Println(l.Rank(e2))

	e3 := l.Add(&item{score: 3})
	fmt.Println(l.Rank(e3))

	e4 := l.Add(&item{score: 2})
	fmt.Println(l.Rank(e4))

	e5 := l.Add(&item{score: 4})
	fmt.Println(l.Rank(e5))

	e6 := l.Add(&item{score: 6})
	fmt.Println(l.Rank(e6))

	fmt.Println()

	// print list after add.
	for e := l.Front(); e != nil; e = e.Next() {
		fmt.Println(e.Value.(*item).score)
	}

	fmt.Println()

	l.Remove(l.Find(6))
	l.Remove(l.Get(4))

	// print list after remoscoree.
	for e := l.Front(); e != nil; e = e.Next() {
		fmt.Println(e.Value.(*item).score)
	}

	// Output:
	// 0
	// 0
	// 1
	// 1
	// 3
	// 5
	//
	// 1
	// 2
	// 3
	// 4
	// 4
	// 6
	//
	// 1
	// 2
	// 3
	// 4
}<|endoftext|>


<commit_before>package tsm1

import (
	"errors"
	"os"
	"reflect"
	"sync"
	"syscall"
	"unsafe"
)

// mmap implementation for Windows
// Based on: https://github.com/edsrzf/mmap-go
// Based on: https://github.com/boltdb/bolt/bolt_windows.go
// Ref: https://groups.google.com/forum/#!topic/golang-nuts/g0nLwQI9www

// We keep this map so that we can get back the original handle from the memory address.
var handleLock sync.Mutex
var handleMap = map[uintptr]syscall.Handle{}
var fileMap = map[uintptr]*os.File{}

func openSharedFile(f *os.File) (file *os.File, err error) {

	var access, createmode, sharemode uint32
	var sa *syscall.SecurityAttributes

	access = syscall.GENERIC_READ
	sharemode = uint32(syscall.FILE_SHARE_READ | syscall.FILE_SHARE_WRITE | syscall.FILE_SHARE_DELETE)
	createmode = syscall.OPEN_EXISTING
	fileName := f.Name()

	pathp, err := syscall.UTF16PtrFromString(fileName)
	if err != nil {
		return nil, err
	}

	h, e := syscall.CreateFile(pathp, access, sharemode, sa, createmode, syscall.FILE_ATTRIBUTE_NORMAL, 0)

	if e != nil {
		return nil, e
	}
	//NewFile does not add finalizer, need to close this manually
	return os.NewFile(uintptr(h), fileName), nil
}

func mmap(f *os.File, offset int64, length int) (out []byte, err error) {
	// Open a file mapping handle.
	sizelo := uint32(length >> 32)
	sizehi := uint32(length) & 0xffffffff

	sharedHandle, errno := openSharedFile(f)
	if errno != nil {
		return nil, os.NewSyscallError("CreateFile", errno)
	}

	h, errno := syscall.CreateFileMapping(syscall.Handle(sharedHandle.Fd()), nil, syscall.PAGE_READONLY, sizelo, sizehi, nil)
	if h == 0 {
		return nil, os.NewSyscallError("CreateFileMapping", errno)
	}

	// Create the memory map.
	addr, errno := syscall.MapViewOfFile(h, syscall.FILE_MAP_READ, 0, 0, uintptr(length))
	if addr == 0 {
		return nil, os.NewSyscallError("MapViewOfFile", errno)
	}

	handleLock.Lock()
	handleMap[addr] = h
	fileMap[addr] = sharedHandle
	handleLock.Unlock()

	// Convert to a byte array.
	hdr := (*reflect.SliceHeader)(unsafe.Pointer(&out))
	hdr.Data = uintptr(unsafe.Pointer(addr))
	hdr.Len = length
	hdr.Cap = length

	return
}

// munmap Windows implementation
// Based on: https://github.com/edsrzf/mmap-go
// Based on: https://github.com/boltdb/bolt/bolt_windows.go
func munmap(b []byte) (err error) {

	handleLock.Lock()
	defer handleLock.Unlock()

	addr := (uintptr)(unsafe.Pointer(&b[0]))
	if err := syscall.UnmapViewOfFile(addr); err != nil {
		return os.NewSyscallError("UnmapViewOfFile", err)
	}

	handle, ok := handleMap[addr]
	if !ok {
		// should be impossible; we would've seen the error above
		return errors.New("unknown base address")
	}
	delete(handleMap, addr)

	e := syscall.CloseHandle(syscall.Handle(handle))
	return os.NewSyscallError("CloseHandle", e)

	file, ok := fileMap[addr]
	if !ok {
		// should be impossible; we would've seen the error above
		return errors.New("unknown base address")
	}
	delete(fileMap, addr)

	e = file.Close()
	return errors.New("Close File" + e.Error())

	return nil
}
<commit_msg>fix munmap bug in the windows<commit_after>package tsm1

import (
	"errors"
	"os"
	"reflect"
	"sync"
	"syscall"
	"unsafe"
)

// mmap implementation for Windows
// Based on: https://github.com/edsrzf/mmap-go
// Based on: https://github.com/boltdb/bolt/bolt_windows.go
// Ref: https://groups.google.com/forum/#!topic/golang-nuts/g0nLwQI9www

// We keep this map so that we can get back the original handle from the memory address.
var handleLock sync.Mutex
var handleMap = map[uintptr]syscall.Handle{}
var fileMap = map[uintptr]*os.File{}

func openSharedFile(f *os.File) (file *os.File, err error) {

	var access, createmode, sharemode uint32
	var sa *syscall.SecurityAttributes

	access = syscall.GENERIC_READ
	sharemode = uint32(syscall.FILE_SHARE_READ | syscall.FILE_SHARE_WRITE | syscall.FILE_SHARE_DELETE)
	createmode = syscall.OPEN_EXISTING
	fileName := f.Name()

	pathp, err := syscall.UTF16PtrFromString(fileName)
	if err != nil {
		return nil, err
	}

	h, e := syscall.CreateFile(pathp, access, sharemode, sa, createmode, syscall.FILE_ATTRIBUTE_NORMAL, 0)

	if e != nil {
		return nil, e
	}
	//NewFile does not add finalizer, need to close this manually
	return os.NewFile(uintptr(h), fileName), nil
}

func mmap(f *os.File, offset int64, length int) (out []byte, err error) {
	// Open a file mapping handle.
	sizelo := uint32(length >> 32)
	sizehi := uint32(length) & 0xffffffff

	sharedHandle, errno := openSharedFile(f)
	if errno != nil {
		return nil, os.NewSyscallError("CreateFile", errno)
	}

	h, errno := syscall.CreateFileMapping(syscall.Handle(sharedHandle.Fd()), nil, syscall.PAGE_READONLY, sizelo, sizehi, nil)
	if h == 0 {
		return nil, os.NewSyscallError("CreateFileMapping", errno)
	}

	// Create the memory map.
	addr, errno := syscall.MapViewOfFile(h, syscall.FILE_MAP_READ, 0, 0, uintptr(length))
	if addr == 0 {
		return nil, os.NewSyscallError("MapViewOfFile", errno)
	}

	handleLock.Lock()
	handleMap[addr] = h
	fileMap[addr] = sharedHandle
	handleLock.Unlock()

	// Convert to a byte array.
	hdr := (*reflect.SliceHeader)(unsafe.Pointer(&out))
	hdr.Data = uintptr(unsafe.Pointer(addr))
	hdr.Len = length
	hdr.Cap = length

	return
}

// munmap Windows implementation
// Based on: https://github.com/edsrzf/mmap-go
// Based on: https://github.com/boltdb/bolt/bolt_windows.go
func munmap(b []byte) (err error) {
	handleLock.Lock()
	defer handleLock.Unlock()

	addr := (uintptr)(unsafe.Pointer(&b[0]))
	if err := syscall.UnmapViewOfFile(addr); err != nil {
		return os.NewSyscallError("UnmapViewOfFile", err)
	}

	handle, ok := handleMap[addr]
	if !ok {
		// should be impossible; we would've seen the error above
		return errors.New("unknown base address")
	}
	delete(handleMap, addr)

	e := syscall.CloseHandle(syscall.Handle(handle))
	if e != nil {
		return os.NewSyscallError("CloseHandle", e)
	}

	file, ok := fileMap[addr]
	if !ok {
		// should be impossible; we would've seen the error above
		return errors.New("unknown base address")
	}
	delete(fileMap, addr)

	e = file.Close()
	if e != nil {
		return errors.New("close file" + e.Error())
	}
	return nil
}
<|endoftext|>


<commit_before>package integration

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"runtime"
	"sort"
	"strings"
	"sync"
	"syscall"
	"testing"

	"github.com/moby/buildkit/util/contentutil"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type Sandbox interface {
	Address() string
	PrintLogs(*testing.T)
	Cmd(...string) *exec.Cmd
	NewRegistry() (string, error)
	Rootless() bool
	Value(string) interface{} // chosen matrix value
}

type Worker interface {
	New(...SandboxOpt) (Sandbox, func() error, error)
	Name() string
}

type SandboxConf struct {
	mirror string
	mv     matrixValue
}

type SandboxOpt func(*SandboxConf)

func WithMirror(h string) SandboxOpt {
	return func(c *SandboxConf) {
		c.mirror = h
	}
}

func withMatrixValues(mv matrixValue) SandboxOpt {
	return func(c *SandboxConf) {
		c.mv = mv
	}
}

type Test func(*testing.T, Sandbox)

var defaultWorkers []Worker

func register(w Worker) {
	defaultWorkers = append(defaultWorkers, w)
}

func List() []Worker {
	return defaultWorkers
}

type TestOpt func(*TestConf)

func WithMatrix(key string, m map[string]interface{}) TestOpt {
	return func(tc *TestConf) {
		if tc.matrix == nil {
			tc.matrix = map[string]map[string]interface{}{}
		}
		tc.matrix[key] = m
	}
}

func WithMirroredImages(m map[string]string) TestOpt {
	return func(tc *TestConf) {
		if tc.mirroredImages == nil {
			tc.mirroredImages = map[string]string{}
		}
		for k, v := range m {
			tc.mirroredImages[k] = v
		}
	}
}

type TestConf struct {
	matrix         map[string]map[string]interface{}
	mirroredImages map[string]string
}

func Run(t *testing.T, testCases []Test, opt ...TestOpt) {
	if testing.Short() {
		t.Skip("skipping in short mode")
	}

	var tc TestConf
	for _, o := range opt {
		o(&tc)
	}

	mirror, cleanup, err := runMirror(t, tc.mirroredImages)
	require.NoError(t, err)

	var mu sync.Mutex
	var count int
	cleanOnComplete := func() func() {
		count++
		return func() {
			mu.Lock()
			count--
			if count == 0 {
				cleanup()
			}
			mu.Unlock()
		}
	}
	defer cleanOnComplete()()

	matrix := prepareValueMatrix(tc)

	for _, br := range List() {
		for _, tc := range testCases {
			for _, mv := range matrix {
				ok := t.Run(getFunctionName(tc)+"/worker="+br.Name()+mv.functionSuffix(), func(t *testing.T) {
					defer cleanOnComplete()()
					sb, close, err := br.New(WithMirror(mirror), withMatrixValues(mv))
					if err != nil {
						if errors.Cause(err) == ErrorRequirements {
							t.Skip(err.Error())
						}
						require.NoError(t, err)
					}
					defer func() {
						assert.NoError(t, close())
						if t.Failed() {
							sb.PrintLogs(t)
						}
					}()
					tc(t, sb)
				})
				require.True(t, ok)
			}
		}
	}
}

func getFunctionName(i interface{}) string {
	fullname := runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
	dot := strings.LastIndex(fullname, ".") + 1
	return strings.Title(fullname[dot:])
}

var localImageCache map[string]map[string]struct{}

func copyImagesLocal(t *testing.T, host string, images map[string]string) error {
	for to, from := range images {
		if localImageCache == nil {
			localImageCache = map[string]map[string]struct{}{}
		}
		if _, ok := localImageCache[host]; !ok {
			localImageCache[host] = map[string]struct{}{}
		}
		if _, ok := localImageCache[host][to]; ok {
			continue
		}
		localImageCache[host][to] = struct{}{}

		desc, provider, err := contentutil.ProviderFromRef(from)
		if err != nil {
			return err
		}
		ingester, err := contentutil.IngesterFromRef(host + "/" + to)
		if err != nil {
			return err
		}
		if err := contentutil.CopyChain(context.TODO(), ingester, provider, desc); err != nil {
			return err
		}
		t.Logf("copied %s to local mirror %s", from, host+"/"+to)
	}
	return nil
}

func OfficialImages(names ...string) map[string]string {
	ns := runtime.GOARCH
	if ns == "arm64" {
		ns = "arm64v8"
	} else if ns != "amd64" && ns != "armhf" {
		ns = "library"
	}
	m := map[string]string{}
	for _, name := range names {
		m["library/"+name] = "docker.io/" + ns + "/" + name
	}
	return m
}

func configWithMirror(mirror string) (string, error) {
	tmpdir, err := ioutil.TempDir("", "bktest_config")
	if err != nil {
		return "", err
	}
	if err := os.Chmod(tmpdir, 0711); err != nil {
		return "", err
	}
	if err := ioutil.WriteFile(filepath.Join(tmpdir, "buildkitd.toml"), []byte(fmt.Sprintf(`
[registry."docker.io"]
mirrors=["%s"]
`, mirror)), 0644); err != nil {
		return "", err
	}
	return tmpdir, nil
}

func runMirror(t *testing.T, mirroredImages map[string]string) (host string, cleanup func() error, err error) {
	mirrorDir := os.Getenv("BUILDKIT_REGISTRY_MIRROR_DIR")

	var f *os.File
	if mirrorDir != "" {
		f, err = os.Create(filepath.Join(mirrorDir, "lock"))
		if err != nil {
			return "", nil, err
		}
		defer func() {
			if err != nil {
				f.Close()
			}
		}()
		if err := syscall.Flock(int(f.Fd()), syscall.LOCK_EX); err != nil {
			return "", nil, err
		}
	}

	mirror, cleanup, err := newRegistry(mirrorDir)
	if err != nil {
		return "", nil, err
	}
	defer func() {
		if err != nil {
			cleanup()
		}
	}()

	if err := copyImagesLocal(t, mirror, mirroredImages); err != nil {
		return "", nil, err
	}

	if mirrorDir != "" {
		if err := syscall.Flock(int(f.Fd()), syscall.LOCK_UN); err != nil {
			return "", nil, err
		}
	}

	return mirror, cleanup, err
}

type matrixValue struct {
	fn     []string
	values map[string]matrixValueChoice
}

func (mv matrixValue) functionSuffix() string {
	if len(mv.fn) == 0 {
		return ""
	}
	sort.Strings(mv.fn)
	sb := &strings.Builder{}
	for _, f := range mv.fn {
		sb.Write([]byte("/" + f + "=" + mv.values[f].name))
	}
	return sb.String()
}

type matrixValueChoice struct {
	name  string
	value interface{}
}

func newMatrixValue(key, name string, v interface{}) matrixValue {
	return matrixValue{
		fn: []string{key},
		values: map[string]matrixValueChoice{
			key: {
				name:  name,
				value: v,
			},
		},
	}
}

func prepareValueMatrix(tc TestConf) []matrixValue {
	m := []matrixValue{}
	for featureName, values := range tc.matrix {
		current := m
		m = []matrixValue{}
		for featureValue, v := range values {
			if len(current) == 0 {
				m = append(m, newMatrixValue(featureName, featureValue, v))
			}
			for _, c := range current {
				vv := newMatrixValue(featureName, featureValue, v)
				vv.fn = append(vv.fn, c.fn...)
				for k, v := range c.values {
					vv.values[k] = v
				}
				m = append(m, vv)
			}
		}
	}
	if len(m) == 0 {
		m = append(m, matrixValue{})
	}
	return m
}
<commit_msg>testutil: avoid overriding cleanup to nil on error<commit_after>package integration

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"runtime"
	"sort"
	"strings"
	"sync"
	"syscall"
	"testing"

	"github.com/moby/buildkit/util/contentutil"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

type Sandbox interface {
	Address() string
	PrintLogs(*testing.T)
	Cmd(...string) *exec.Cmd
	NewRegistry() (string, error)
	Rootless() bool
	Value(string) interface{} // chosen matrix value
}

type Worker interface {
	New(...SandboxOpt) (Sandbox, func() error, error)
	Name() string
}

type SandboxConf struct {
	mirror string
	mv     matrixValue
}

type SandboxOpt func(*SandboxConf)

func WithMirror(h string) SandboxOpt {
	return func(c *SandboxConf) {
		c.mirror = h
	}
}

func withMatrixValues(mv matrixValue) SandboxOpt {
	return func(c *SandboxConf) {
		c.mv = mv
	}
}

type Test func(*testing.T, Sandbox)

var defaultWorkers []Worker

func register(w Worker) {
	defaultWorkers = append(defaultWorkers, w)
}

func List() []Worker {
	return defaultWorkers
}

type TestOpt func(*TestConf)

func WithMatrix(key string, m map[string]interface{}) TestOpt {
	return func(tc *TestConf) {
		if tc.matrix == nil {
			tc.matrix = map[string]map[string]interface{}{}
		}
		tc.matrix[key] = m
	}
}

func WithMirroredImages(m map[string]string) TestOpt {
	return func(tc *TestConf) {
		if tc.mirroredImages == nil {
			tc.mirroredImages = map[string]string{}
		}
		for k, v := range m {
			tc.mirroredImages[k] = v
		}
	}
}

type TestConf struct {
	matrix         map[string]map[string]interface{}
	mirroredImages map[string]string
}

func Run(t *testing.T, testCases []Test, opt ...TestOpt) {
	if testing.Short() {
		t.Skip("skipping in short mode")
	}

	var tc TestConf
	for _, o := range opt {
		o(&tc)
	}

	mirror, cleanup, err := runMirror(t, tc.mirroredImages)
	require.NoError(t, err)

	var mu sync.Mutex
	var count int
	cleanOnComplete := func() func() {
		count++
		return func() {
			mu.Lock()
			count--
			if count == 0 {
				cleanup()
			}
			mu.Unlock()
		}
	}
	defer cleanOnComplete()()

	matrix := prepareValueMatrix(tc)

	for _, br := range List() {
		for _, tc := range testCases {
			for _, mv := range matrix {
				ok := t.Run(getFunctionName(tc)+"/worker="+br.Name()+mv.functionSuffix(), func(t *testing.T) {
					defer cleanOnComplete()()
					sb, close, err := br.New(WithMirror(mirror), withMatrixValues(mv))
					if err != nil {
						if errors.Cause(err) == ErrorRequirements {
							t.Skip(err.Error())
						}
						require.NoError(t, err)
					}
					defer func() {
						assert.NoError(t, close())
						if t.Failed() {
							sb.PrintLogs(t)
						}
					}()
					tc(t, sb)
				})
				require.True(t, ok)
			}
		}
	}
}

func getFunctionName(i interface{}) string {
	fullname := runtime.FuncForPC(reflect.ValueOf(i).Pointer()).Name()
	dot := strings.LastIndex(fullname, ".") + 1
	return strings.Title(fullname[dot:])
}

var localImageCache map[string]map[string]struct{}

func copyImagesLocal(t *testing.T, host string, images map[string]string) error {
	for to, from := range images {
		if localImageCache == nil {
			localImageCache = map[string]map[string]struct{}{}
		}
		if _, ok := localImageCache[host]; !ok {
			localImageCache[host] = map[string]struct{}{}
		}
		if _, ok := localImageCache[host][to]; ok {
			continue
		}
		localImageCache[host][to] = struct{}{}

		desc, provider, err := contentutil.ProviderFromRef(from)
		if err != nil {
			return err
		}
		ingester, err := contentutil.IngesterFromRef(host + "/" + to)
		if err != nil {
			return err
		}
		if err := contentutil.CopyChain(context.TODO(), ingester, provider, desc); err != nil {
			return err
		}
		t.Logf("copied %s to local mirror %s", from, host+"/"+to)
	}
	return nil
}

func OfficialImages(names ...string) map[string]string {
	ns := runtime.GOARCH
	if ns == "arm64" {
		ns = "arm64v8"
	} else if ns != "amd64" && ns != "armhf" {
		ns = "library"
	}
	m := map[string]string{}
	for _, name := range names {
		m["library/"+name] = "docker.io/" + ns + "/" + name
	}
	return m
}

func configWithMirror(mirror string) (string, error) {
	tmpdir, err := ioutil.TempDir("", "bktest_config")
	if err != nil {
		return "", err
	}
	if err := os.Chmod(tmpdir, 0711); err != nil {
		return "", err
	}
	if err := ioutil.WriteFile(filepath.Join(tmpdir, "buildkitd.toml"), []byte(fmt.Sprintf(`
[registry."docker.io"]
mirrors=["%s"]
`, mirror)), 0644); err != nil {
		return "", err
	}
	return tmpdir, nil
}

func runMirror(t *testing.T, mirroredImages map[string]string) (host string, _ func() error, err error) {
	mirrorDir := os.Getenv("BUILDKIT_REGISTRY_MIRROR_DIR")

	var f *os.File
	if mirrorDir != "" {
		f, err = os.Create(filepath.Join(mirrorDir, "lock"))
		if err != nil {
			return "", nil, err
		}
		defer func() {
			if err != nil {
				f.Close()
			}
		}()
		if err := syscall.Flock(int(f.Fd()), syscall.LOCK_EX); err != nil {
			return "", nil, err
		}
	}

	mirror, cleanup, err := newRegistry(mirrorDir)
	if err != nil {
		return "", nil, err
	}
	defer func() {
		if err != nil {
			cleanup()
		}
	}()

	if err := copyImagesLocal(t, mirror, mirroredImages); err != nil {
		return "", nil, err
	}

	if mirrorDir != "" {
		if err := syscall.Flock(int(f.Fd()), syscall.LOCK_UN); err != nil {
			return "", nil, err
		}
	}

	return mirror, cleanup, err
}

type matrixValue struct {
	fn     []string
	values map[string]matrixValueChoice
}

func (mv matrixValue) functionSuffix() string {
	if len(mv.fn) == 0 {
		return ""
	}
	sort.Strings(mv.fn)
	sb := &strings.Builder{}
	for _, f := range mv.fn {
		sb.Write([]byte("/" + f + "=" + mv.values[f].name))
	}
	return sb.String()
}

type matrixValueChoice struct {
	name  string
	value interface{}
}

func newMatrixValue(key, name string, v interface{}) matrixValue {
	return matrixValue{
		fn: []string{key},
		values: map[string]matrixValueChoice{
			key: {
				name:  name,
				value: v,
			},
		},
	}
}

func prepareValueMatrix(tc TestConf) []matrixValue {
	m := []matrixValue{}
	for featureName, values := range tc.matrix {
		current := m
		m = []matrixValue{}
		for featureValue, v := range values {
			if len(current) == 0 {
				m = append(m, newMatrixValue(featureName, featureValue, v))
			}
			for _, c := range current {
				vv := newMatrixValue(featureName, featureValue, v)
				vv.fn = append(vv.fn, c.fn...)
				for k, v := range c.values {
					vv.values[k] = v
				}
				m = append(m, vv)
			}
		}
	}
	if len(m) == 0 {
		m = append(m, matrixValue{})
	}
	return m
}
<|endoftext|>


<commit_before>package rpcd

import (
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/Symantec/Dominator/lib/hash"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/proto/objectserver"
)

var exclusive sync.RWMutex

func (objSrv *srpcType) GetObjects(conn *srpc.Conn, decoder srpc.Decoder,
	encoder srpc.Encoder) error {
	defer conn.Flush()
	var request objectserver.GetObjectsRequest
	var response objectserver.GetObjectsResponse
	if request.Exclusive {
		exclusive.Lock()
		defer exclusive.Unlock()
	} else {
		exclusive.RLock()
		defer exclusive.RUnlock()
		objSrv.getSemaphore <- true
		defer releaseSemaphore(objSrv.getSemaphore)
	}
	var err error
	if err = decoder.Decode(&request); err != nil {
		response.ResponseString = err.Error()
		return encoder.Encode(response)
	}
	response.ObjectSizes, err = objSrv.objectServer.CheckObjects(request.Hashes)
	if err != nil {
		response.ResponseString = err.Error()
		return encoder.Encode(response)
	}
	// First a quick check for existence. If any objects missing, fail request.
	var firstMissingObject *hash.Hash
	numMissingObjects := 0
	for index, hashVal := range request.Hashes {
		if response.ObjectSizes[index] < 1 {
			firstMissingObject = &hashVal
			numMissingObjects++
		}
	}
	if firstMissingObject != nil {
		if numMissingObjects == 1 {
			response.ResponseString = fmt.Sprintf("unknown object: %x",
				*firstMissingObject)
		} else {
			response.ResponseString = fmt.Sprintf(
				"first of %d unknown objects: %x", numMissingObjects,
				*firstMissingObject)
		}
		return encoder.Encode(response)
	}
	objectsReader, err := objSrv.objectServer.GetObjects(request.Hashes)
	if err != nil {
		response.ResponseString = err.Error()
		return encoder.Encode(response)
	}
	defer objectsReader.Close()
	if err := encoder.Encode(response); err != nil {
		return err
	}
	conn.Flush()
	for _, hashVal := range request.Hashes {
		length, reader, err := objectsReader.NextObject()
		if err != nil {
			objSrv.logger.Println(err)
			return err
		}
		nCopied, err := io.Copy(conn.Writer, reader)
		reader.Close()
		if err != nil {
			objSrv.logger.Printf("Error copying: %s\n", err)
			return err
		}
		if nCopied != int64(length) {
			txt := fmt.Sprintf("Expected length: %d, got: %d for: %x",
				length, nCopied, hashVal)
			objSrv.logger.Printf(txt)
			return errors.New(txt)
		}
	}
	objSrv.logger.Debugf(0, "GetObjects() sent: %d objects\n",
		len(request.Hashes))
	return nil
}

func releaseSemaphore(semaphore <-chan bool) {
	<-semaphore
}
<commit_msg>Re-use same buffer in ObjectServer.GetObjects SRPC method.<commit_after>package rpcd

import (
	"errors"
	"fmt"
	"io"
	"sync"

	"github.com/Symantec/Dominator/lib/hash"
	"github.com/Symantec/Dominator/lib/srpc"
	"github.com/Symantec/Dominator/proto/objectserver"
)

var exclusive sync.RWMutex

func (objSrv *srpcType) GetObjects(conn *srpc.Conn, decoder srpc.Decoder,
	encoder srpc.Encoder) error {
	defer conn.Flush()
	var request objectserver.GetObjectsRequest
	var response objectserver.GetObjectsResponse
	if request.Exclusive {
		exclusive.Lock()
		defer exclusive.Unlock()
	} else {
		exclusive.RLock()
		defer exclusive.RUnlock()
		objSrv.getSemaphore <- true
		defer releaseSemaphore(objSrv.getSemaphore)
	}
	var err error
	if err = decoder.Decode(&request); err != nil {
		response.ResponseString = err.Error()
		return encoder.Encode(response)
	}
	response.ObjectSizes, err = objSrv.objectServer.CheckObjects(request.Hashes)
	if err != nil {
		response.ResponseString = err.Error()
		return encoder.Encode(response)
	}
	// First a quick check for existence. If any objects missing, fail request.
	var firstMissingObject *hash.Hash
	numMissingObjects := 0
	for index, hashVal := range request.Hashes {
		if response.ObjectSizes[index] < 1 {
			firstMissingObject = &hashVal
			numMissingObjects++
		}
	}
	if firstMissingObject != nil {
		if numMissingObjects == 1 {
			response.ResponseString = fmt.Sprintf("unknown object: %x",
				*firstMissingObject)
		} else {
			response.ResponseString = fmt.Sprintf(
				"first of %d unknown objects: %x", numMissingObjects,
				*firstMissingObject)
		}
		return encoder.Encode(response)
	}
	objectsReader, err := objSrv.objectServer.GetObjects(request.Hashes)
	if err != nil {
		response.ResponseString = err.Error()
		return encoder.Encode(response)
	}
	defer objectsReader.Close()
	if err := encoder.Encode(response); err != nil {
		return err
	}
	conn.Flush()
	buffer := make([]byte, 32<<10)
	for _, hashVal := range request.Hashes {
		length, reader, err := objectsReader.NextObject()
		if err != nil {
			objSrv.logger.Println(err)
			return err
		}
		nCopied, err := io.CopyBuffer(conn, reader, buffer)
		reader.Close()
		if err != nil {
			objSrv.logger.Printf("Error copying: %s\n", err)
			return err
		}
		if nCopied != int64(length) {
			txt := fmt.Sprintf("Expected length: %d, got: %d for: %x",
				length, nCopied, hashVal)
			objSrv.logger.Printf(txt)
			return errors.New(txt)
		}
	}
	objSrv.logger.Debugf(0, "GetObjects() sent: %d objects\n",
		len(request.Hashes))
	return nil
}

func releaseSemaphore(semaphore <-chan bool) {
	<-semaphore
}
<|endoftext|>


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

package api

// TODO(ericsnow) Eliminate the dependence on apiserver if possible.

import (
	"strings"
	"time"

	"github.com/juju/errors"
	"github.com/juju/names"
	"gopkg.in/juju/charm.v6-unstable"
	charmresource "gopkg.in/juju/charm.v6-unstable/resource"

	"github.com/juju/juju/apiserver/params"
)

// ListResourcesArgs are the arguments for the ListResources endpoint.
type ListResourcesArgs params.Entities

// NewListResourcesArgs returns the arguments for the ListResources endpoint.
func NewListResourcesArgs(services []string) (ListResourcesArgs, error) {
	var args ListResourcesArgs
	var errs []error
	for _, service := range services {
		if !names.IsValidService(service) {
			err := errors.Errorf("invalid service %q", service)
			errs = append(errs, err)
			continue
		}
		args.Entities = append(args.Entities, params.Entity{
			Tag: names.NewServiceTag(service).String(),
		})
	}
	if err := resolveErrors(errs); err != nil {
		return args, errors.Trace(err)
	}
	return args, nil
}

// AddPendingResourcesArgs holds the arguments to the AddPendingResources
// API endpoint.
type AddPendingResourcesArgs struct {
	params.Entity
	params.AddCharmWithAuthorization

	// Resources is the list of resources to add as pending.
	Resources []CharmResource
}

// NewAddPendingResourcesArgs returns the arguments for the
// AddPendingResources API endpoint.
func NewAddPendingResourcesArgs(serviceID string, cURL *charm.URL, resources []charmresource.Resource) (AddPendingResourcesArgs, error) {
	var args AddPendingResourcesArgs

	if !names.IsValidService(serviceID) {
		return args, errors.Errorf("invalid service %q", serviceID)
	}
	tag := names.NewServiceTag(serviceID).String()

	var apiResources []CharmResource
	for _, res := range resources {
		if err := res.Validate(); err != nil {
			return args, errors.Trace(err)
		}
		apiRes := CharmResource2API(res)
		apiResources = append(apiResources, apiRes)
	}
	args.Tag = tag
	args.Resources = apiResources
	args.URL = cURL.String()
	// TODO(ericsnow) Set args.CharmStoreMacaroon.
	return args, nil
}

// AddPendingResourcesResult holds the result of the AddPendingResources
// API endpoint.
type AddPendingResourcesResult struct {
	params.ErrorResult

	// PendingIDs holds the "pending ID" for each of the requested
	// resources.
	PendingIDs []string
}

// ResourcesResults holds the resources that result
// from a bulk API call.
type ResourcesResults struct {
	// Results is the list of resource results.
	Results []ResourcesResult
}

// ResourcesResult holds the resources that result from an API call
// for a single service.
type ResourcesResult struct {
	params.ErrorResult

	// Resources is the list of resources for the service.
	Resources []Resource

	// CharmStoreResources is the list of resources associated with the charm in
	// the charmstore.
	CharmStoreResources []CharmResource

	// UnitResources contains a list of the resources for each unit in the
	// service.
	UnitResources []UnitResources
}

// A UnitResources contains a list of the resources the unit defined by Entity.
type UnitResources struct {
	params.Entity

	// Resources is a list of resources for the unit.
	Resources []Resource
}

// UploadResult is the response from an upload request.
type UploadResult struct {
	params.ErrorResult

	// Resource describes the resource that was stored in the model.
	Resource Resource
}

// Resource contains info about a Resource.
type Resource struct {
	CharmResource

	// ID uniquely identifies a resource-service pair within the model.
	// Note that the model ignores pending resources (those with a
	// pending ID) except for in a few clearly pending-related places.
	ID string

	// PendingID identifies that this resource is pending and
	// distinguishes it from other pending resources with the same model
	// ID (and from the active resource).
	PendingID string

	// ServiceID identifies the service for the resource.
	ServiceID string

	// Username is the ID of the user that added the revision
	// to the model (whether implicitly or explicitly).
	Username string `json:"username"`

	// Timestamp indicates when the resource was added to the model.
	Timestamp time.Time `json:"timestamp"`
}

// CharmResource contains the definition for a resource.
type CharmResource struct {
	// Name identifies the resource.
	Name string `json:"name"`

	// Type is the name of the resource type.
	Type string `json:"type"`

	// Path is where the resource will be stored.
	Path string `json:"path"`

	// Description contains user-facing info about the resource.
	Description string `json:"description,omitempty"`

	// Origin is where the resource will come from.
	Origin string `json:"origin"`

	// Revision is the revision, if applicable.
	Revision int `json:"revision"`

	// Fingerprint is the SHA-384 checksum for the resource blob.
	Fingerprint []byte `json:"fingerprint"`

	// Size is the size of the resource, in bytes.
	Size int64
}

func resolveErrors(errs []error) error {
	switch len(errs) {
	case 0:
		return nil
	case 1:
		return errs[0]
	default:
		msgs := make([]string, len(errs))
		for i, err := range errs {
			msgs[i] = err.Error()
		}
		return errors.New(strings.Join(msgs, "\n"))
	}
}
<commit_msg>Only include the URL arg if there is a URL.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package api

// TODO(ericsnow) Eliminate the dependence on apiserver if possible.

import (
	"strings"
	"time"

	"github.com/juju/errors"
	"github.com/juju/names"
	"gopkg.in/juju/charm.v6-unstable"
	charmresource "gopkg.in/juju/charm.v6-unstable/resource"

	"github.com/juju/juju/apiserver/params"
)

// ListResourcesArgs are the arguments for the ListResources endpoint.
type ListResourcesArgs params.Entities

// NewListResourcesArgs returns the arguments for the ListResources endpoint.
func NewListResourcesArgs(services []string) (ListResourcesArgs, error) {
	var args ListResourcesArgs
	var errs []error
	for _, service := range services {
		if !names.IsValidService(service) {
			err := errors.Errorf("invalid service %q", service)
			errs = append(errs, err)
			continue
		}
		args.Entities = append(args.Entities, params.Entity{
			Tag: names.NewServiceTag(service).String(),
		})
	}
	if err := resolveErrors(errs); err != nil {
		return args, errors.Trace(err)
	}
	return args, nil
}

// AddPendingResourcesArgs holds the arguments to the AddPendingResources
// API endpoint.
type AddPendingResourcesArgs struct {
	params.Entity
	params.AddCharmWithAuthorization

	// Resources is the list of resources to add as pending.
	Resources []CharmResource
}

// NewAddPendingResourcesArgs returns the arguments for the
// AddPendingResources API endpoint.
func NewAddPendingResourcesArgs(serviceID string, cURL *charm.URL, resources []charmresource.Resource) (AddPendingResourcesArgs, error) {
	var args AddPendingResourcesArgs

	if !names.IsValidService(serviceID) {
		return args, errors.Errorf("invalid service %q", serviceID)
	}
	tag := names.NewServiceTag(serviceID).String()

	var apiResources []CharmResource
	for _, res := range resources {
		if err := res.Validate(); err != nil {
			return args, errors.Trace(err)
		}
		apiRes := CharmResource2API(res)
		apiResources = append(apiResources, apiRes)
	}
	args.Tag = tag
	args.Resources = apiResources
	if cURL != nil {
		args.URL = cURL.String()
		// TODO(ericsnow) Set args.CharmStoreMacaroon.
	}
	return args, nil
}

// AddPendingResourcesResult holds the result of the AddPendingResources
// API endpoint.
type AddPendingResourcesResult struct {
	params.ErrorResult

	// PendingIDs holds the "pending ID" for each of the requested
	// resources.
	PendingIDs []string
}

// ResourcesResults holds the resources that result
// from a bulk API call.
type ResourcesResults struct {
	// Results is the list of resource results.
	Results []ResourcesResult
}

// ResourcesResult holds the resources that result from an API call
// for a single service.
type ResourcesResult struct {
	params.ErrorResult

	// Resources is the list of resources for the service.
	Resources []Resource

	// CharmStoreResources is the list of resources associated with the charm in
	// the charmstore.
	CharmStoreResources []CharmResource

	// UnitResources contains a list of the resources for each unit in the
	// service.
	UnitResources []UnitResources
}

// A UnitResources contains a list of the resources the unit defined by Entity.
type UnitResources struct {
	params.Entity

	// Resources is a list of resources for the unit.
	Resources []Resource
}

// UploadResult is the response from an upload request.
type UploadResult struct {
	params.ErrorResult

	// Resource describes the resource that was stored in the model.
	Resource Resource
}

// Resource contains info about a Resource.
type Resource struct {
	CharmResource

	// ID uniquely identifies a resource-service pair within the model.
	// Note that the model ignores pending resources (those with a
	// pending ID) except for in a few clearly pending-related places.
	ID string

	// PendingID identifies that this resource is pending and
	// distinguishes it from other pending resources with the same model
	// ID (and from the active resource).
	PendingID string

	// ServiceID identifies the service for the resource.
	ServiceID string

	// Username is the ID of the user that added the revision
	// to the model (whether implicitly or explicitly).
	Username string `json:"username"`

	// Timestamp indicates when the resource was added to the model.
	Timestamp time.Time `json:"timestamp"`
}

// CharmResource contains the definition for a resource.
type CharmResource struct {
	// Name identifies the resource.
	Name string `json:"name"`

	// Type is the name of the resource type.
	Type string `json:"type"`

	// Path is where the resource will be stored.
	Path string `json:"path"`

	// Description contains user-facing info about the resource.
	Description string `json:"description,omitempty"`

	// Origin is where the resource will come from.
	Origin string `json:"origin"`

	// Revision is the revision, if applicable.
	Revision int `json:"revision"`

	// Fingerprint is the SHA-384 checksum for the resource blob.
	Fingerprint []byte `json:"fingerprint"`

	// Size is the size of the resource, in bytes.
	Size int64
}

func resolveErrors(errs []error) error {
	switch len(errs) {
	case 0:
		return nil
	case 1:
		return errs[0]
	default:
		msgs := make([]string, len(errs))
		for i, err := range errs {
			msgs[i] = err.Error()
		}
		return errors.New(strings.Join(msgs, "\n"))
	}
}
<|endoftext|>


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

package cmd

import (
	"fmt"
	"io"

	"github.com/juju/cmd"
	"github.com/juju/errors"
	"launchpad.net/gnuflag"

	"github.com/juju/juju/cmd/envcmd"
	"github.com/juju/juju/resource"
)

// ShowAPI has the API methods needed by ShowCommand.
type ShowAPI interface {
	// ListSpecs lists the resource specs for each of the given services.
	ListSpecs(services ...string) ([]resource.SpecsResult, error)

	// Close closes the client.
	Close() error
}

// ShowCommand implements the show-resources command.
type ShowCommand struct {
	envcmd.EnvCommandBase
	out       cmd.Output
	serviceID string

	newAPIClient func(c *ShowCommand) (ShowAPI, error)
}

// NewShowCommand returns a new command that lists resources defined
// by a charm.
func NewShowCommand(newAPIClient func(c *ShowCommand) (ShowAPI, error)) *ShowCommand {
	cmd := &ShowCommand{
		newAPIClient: newAPIClient,
	}
	return cmd
}

var showDoc = `
This command will report the resources defined by a charm.

The resources are looked up in the service's charm metadata.
`

func (c *ShowCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "show-resources",
		Args:    "service-id",
		Purpose: "display the charm-defined resources for a service",
		Doc:     showDoc,
	}
}

func (c *ShowCommand) SetFlags(f *gnuflag.FlagSet) {
	defaultFormat := "tabular"
	c.out.AddFlags(f, defaultFormat, map[string]cmd.Formatter{
		"tabular": FormatTabular,
		"yaml":    cmd.FormatYaml,
		"json":    cmd.FormatJson,
	})
}

func (c *ShowCommand) Init(args []string) error {
	if len(args) == 0 {
		return errors.New("missing service ID")
	}
	c.serviceID = args[0]

	if err := cmd.CheckEmpty(args[1:]); err != nil {
		return errors.Trace(err)
	}

	return nil
}

// TODO(ericsnow) Move this to a common place, like cmd/envcmd?
const connectionError = `Unable to connect to environment %q.
Please check your credentials or use 'juju bootstrap' to create a new environment.

Error details:
%v
`

func (c *ShowCommand) Run(ctx *cmd.Context) error {
	apiclient, err := c.newAPIClient(c)
	if err != nil {
		return fmt.Errorf(connectionError, c.ConnectionName(), err)
	}
	defer apiclient.Close()

	results, err := apiclient.ListSpecs(c.serviceID)
	if err != nil {
		return errors.Trace(err)
	}

	var specs []resource.Spec
	if len(results) == 1 {
		// TODO(ericsnow) Handle results[0].Error?
		specs = results[0].Specs
	}

	// Note that we do not worry about c.CompatVersion for show-resources...
	formatter := newSpecListFormatter(specs)
	formatted := formatter.format()
	return c.out.Write(ctx, formatted)
}
<commit_msg>Add some doc comments.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package cmd

import (
	"fmt"
	"io"

	"github.com/juju/cmd"
	"github.com/juju/errors"
	"launchpad.net/gnuflag"

	"github.com/juju/juju/cmd/envcmd"
	"github.com/juju/juju/resource"
)

// ShowAPI has the API methods needed by ShowCommand.
type ShowAPI interface {
	// ListSpecs lists the resource specs for each of the given services.
	ListSpecs(services ...string) ([]resource.SpecsResult, error)

	// Close closes the client.
	Close() error
}

// ShowCommand implements the show-resources command.
type ShowCommand struct {
	envcmd.EnvCommandBase
	out       cmd.Output
	serviceID string

	newAPIClient func(c *ShowCommand) (ShowAPI, error)
}

// NewShowCommand returns a new command that lists resources defined
// by a charm.
func NewShowCommand(newAPIClient func(c *ShowCommand) (ShowAPI, error)) *ShowCommand {
	cmd := &ShowCommand{
		newAPIClient: newAPIClient,
	}
	return cmd
}

var showDoc = `
This command will report the resources defined by a charm.

The resources are looked up in the service's charm metadata.
`

// Info implements cmd.Command.
func (c *ShowCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "show-resources",
		Args:    "service-id",
		Purpose: "display the charm-defined resources for a service",
		Doc:     showDoc,
	}
}

// SetFlags implements cmd.Command.
func (c *ShowCommand) SetFlags(f *gnuflag.FlagSet) {
	defaultFormat := "tabular"
	c.out.AddFlags(f, defaultFormat, map[string]cmd.Formatter{
		"tabular": FormatTabular,
		"yaml":    cmd.FormatYaml,
		"json":    cmd.FormatJson,
	})
}

// Init implements cmd.Command.
func (c *ShowCommand) Init(args []string) error {
	if len(args) == 0 {
		return errors.New("missing service ID")
	}
	c.serviceID = args[0]

	if err := cmd.CheckEmpty(args[1:]); err != nil {
		return errors.Trace(err)
	}

	return nil
}

// TODO(ericsnow) Move this to a common place, like cmd/envcmd?
const connectionError = `Unable to connect to environment %q.
Please check your credentials or use 'juju bootstrap' to create a new environment.

Error details:
%v
`

// Run implements cmd.Command.
func (c *ShowCommand) Run(ctx *cmd.Context) error {
	apiclient, err := c.newAPIClient(c)
	if err != nil {
		return fmt.Errorf(connectionError, c.ConnectionName(), err)
	}
	defer apiclient.Close()

	results, err := apiclient.ListSpecs(c.serviceID)
	if err != nil {
		return errors.Trace(err)
	}

	var specs []resource.Spec
	if len(results) == 1 {
		// TODO(ericsnow) Handle results[0].Error?
		specs = results[0].Specs
	}

	// Note that we do not worry about c.CompatVersion for show-resources...
	formatter := newSpecListFormatter(specs)
	formatted := formatter.format()
	return c.out.Write(ctx, formatted)
}
<|endoftext|>


<commit_before>package main

const stmtGetDeployment = `
SELECT scic.check_instance_config_id,
       scic.status,
       scic.next_status,
       scic.deployment_details
FROM   soma.check_instances sci
JOIN   soma.check_instance_configurations scic
ON     sci.check_instance_id = scic.check_instance_id
AND    sci.current_instance_config_id = scic.check_instance_config_id
WHERE  sci.check_instance_id = $1::uuid
AND    (  sci.status = 'awaiting_rollout'
       OR sci.status = 'rollout_in_progress'
	   OR sci.status = 'active'
	   OR sci.status = 'rollout_failed'
	   OR sci.status = 'awaiting_deprovision'
	   OR sci.status = 'deprovision_in_progress'
       OR sci.status = 'deprovision_failed' );`

const stmtUpdateDeployment = `
UPDATE soma.check_instance_configurations
SET    status = $1::varchar,
       next_status = $2::varchar
WHERE  check_instance_config_id = $3::uuid;`

const stmtDeploymentStatus = `
SELECT scic.check_instance_config_id,
       scic.status,
	   scic.next_status
FROM   soma.check_instances sci
JOIN   soma.check_instance_configurations scic
ON     sci.check_instance_id = scic.check_instance_id
AND    sci.current_instance_config_id = scic.check_instance_config_id
WHERE  sci.check_instance_id = $1::uuid;`

const stmtActivateDeployment = `
UPDATE soma.check_instance_configurations
SET    status = $1::varchar,
       next_status = $2::varchar,
	   activated_at = $3::timestamptz,
WHERE  check_instance_config_id = $4::uuid;`

const stmtGetDeploymentList = `
SELECT sci.instance_id
FROM   soma.monitoring_systems sms
JOIN   soma.check_instance_configurations scic
ON     sms.monitoring_id = scic.monitoring_id
JOIN   soma.check_instances sci
ON     scic.check_instance_id = sci.check_instance_id
AND    scic.check_instance_config_id = sci.current_instance_config_id
WHERE  sms.monitoring_id = $1::uuid
AND    sci.update_available
AND    (  scic.status = 'awaiting_rollout'
       OR scic.status = 'awaiting_deprovision' );`

const stmtGetAllDeploymentList = `
SELECT sci.instance_id
FROM   soma.monitoring_systems sms
JOIN   soma.check_instance_configurations scic
ON     sms.monitoring_id = scic.monitoring_id
JOIN   soma.check_instances sci
ON     scic.check_instance_id = sci.check_instance_id
AND    scic.check_instance_config_id = sci.current_instance_config_id
WHERE  sms.monitoring_id = $1::uuid
AND    (  scic.status = 'awaiting_rollout'
	   OR scic.status = 'rollout_in_progress'
       OR scic.status = 'awaiting_deprovision'
	   OR scic.status = 'deprovision_in_progress');`

const stmtDeployClearFlag = `
UPDATE soma.check_instances
SET    update_available = 'false'::boolean
WHERE  check_instance_id = $1::uuid;`

const stmtGetInstancesForNode = `
SELECT sci.check_instance_id
FROM   soma.nodes sn
JOIN   soma.checks sc
ON     sn.node_id = sc.object_id
JOIN   soma.monitoring_capabilities smc
ON     sc.capability_id = smc.monitoring_capability
JOIN   soma.check_instances sci
ON     sc.check_id = sci.check_id
WHERE  sn.node_asset_id = $1::numeric
AND    sc.object_type = 'node'
AND    smc.capability_view = 'local'
AND    smc.capability_monitoring = $2::uuid;`

const stmtGetLastInstanceVersion = `
SELECT deployment_details
FROM   soma.check_instance_configurations
WHERE  check_instance_id = $1::uuid
AND    (   status != 'deprovisioned'
       AND status != 'awaiting_deletion')
ORDER  BY version DESC
LIMIT  1;`

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Fix SQL statement syntax errors<commit_after>package main

const stmtGetDeployment = `
SELECT scic.check_instance_config_id,
       scic.status,
       scic.next_status,
       scic.deployment_details
FROM   soma.check_instances sci
JOIN   soma.check_instance_configurations scic
ON     sci.check_instance_id = scic.check_instance_id
AND    sci.current_instance_config_id = scic.check_instance_config_id
WHERE  sci.check_instance_id = $1::uuid
AND    (  scic.status = 'awaiting_rollout'
       OR scic.status = 'rollout_in_progress'
	   OR scic.status = 'active'
	   OR scic.status = 'rollout_failed'
	   OR scic.status = 'awaiting_deprovision'
	   OR scic.status = 'deprovision_in_progress'
       OR scic.status = 'deprovision_failed' );`

const stmtUpdateDeployment = `
UPDATE soma.check_instance_configurations
SET    status = $1::varchar,
       next_status = $2::varchar
WHERE  check_instance_config_id = $3::uuid;`

const stmtDeploymentStatus = `
SELECT scic.check_instance_config_id,
       scic.status,
	   scic.next_status
FROM   soma.check_instances sci
JOIN   soma.check_instance_configurations scic
ON     sci.check_instance_id = scic.check_instance_id
AND    sci.current_instance_config_id = scic.check_instance_config_id
WHERE  sci.check_instance_id = $1::uuid;`

const stmtActivateDeployment = `
UPDATE soma.check_instance_configurations
SET    status = $1::varchar,
       next_status = $2::varchar,
	   activated_at = $3::timestamptz
WHERE  check_instance_config_id = $4::uuid;`

const stmtGetDeploymentList = `
SELECT sci.check_instance_id
FROM   soma.monitoring_systems sms
JOIN   soma.check_instance_configurations scic
ON     sms.monitoring_id = scic.monitoring_id
JOIN   soma.check_instances sci
ON     scic.check_instance_id = sci.check_instance_id
AND    scic.check_instance_config_id = sci.current_instance_config_id
WHERE  sms.monitoring_id = $1::uuid
AND    sci.update_available
AND    (  scic.status = 'awaiting_rollout'
       OR scic.status = 'awaiting_deprovision' );`

const stmtGetAllDeploymentList = `
SELECT sci.check_instance_id
FROM   soma.monitoring_systems sms
JOIN   soma.check_instance_configurations scic
ON     sms.monitoring_id = scic.monitoring_id
JOIN   soma.check_instances sci
ON     scic.check_instance_id = sci.check_instance_id
AND    scic.check_instance_config_id = sci.current_instance_config_id
WHERE  sms.monitoring_id = $1::uuid
AND    (  scic.status = 'awaiting_rollout'
	   OR scic.status = 'rollout_in_progress'
       OR scic.status = 'awaiting_deprovision'
	   OR scic.status = 'deprovision_in_progress');`

const stmtDeployClearFlag = `
UPDATE soma.check_instances
SET    update_available = 'false'::boolean
WHERE  check_instance_id = $1::uuid;`

const stmtGetInstancesForNode = `
SELECT sci.check_instance_id
FROM   soma.nodes sn
JOIN   soma.checks sc
ON     sn.node_id = sc.object_id
JOIN   soma.monitoring_capabilities smc
ON     sc.capability_id = smc.capability_id
JOIN   soma.check_instances sci
ON     sc.check_id = sci.check_id
WHERE  sn.node_asset_id = $1::numeric
AND    sc.object_type = 'node'
AND    smc.capability_view = 'local'
AND    smc.capability_monitoring = $2::uuid;`

const stmtGetLastInstanceVersion = `
SELECT deployment_details
FROM   soma.check_instance_configurations
WHERE  check_instance_id = $1::uuid
AND    (   status != 'deprovisioned'
       AND status != 'awaiting_deletion')
ORDER  BY version DESC
LIMIT  1;`

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package utee

import (
	"container/heap"
	"sync"
	"time"
)

// An Item is something we manage in a priority queue.
type pqItem struct {
	key   interface{}
	value interface{}
	ttl   int64 //unix timestamp , in second
	index int
	dead  bool //mark as dead,
}

// A priorityQueue implements heap.Interface and holds Items.
type priorityQueue []*pqItem

func (pq priorityQueue) Len() int { return len(pq) }

func (pq priorityQueue) Less(i, j int) bool {
	return pq[i].ttl < pq[j].ttl
}

func (pq priorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}

func (pq *priorityQueue) Push(x interface{}) {
	n := len(*pq)
	item := x.(*pqItem)
	item.index = n
	*pq = append(*pq, item)
}

func (pq *priorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	item.index = -1 // for safety
	*pq = old[0 : n-1]
	return item
}

type TimerCache struct {
	sync.RWMutex

	q   priorityQueue
	m   map[interface{}]*pqItem
	ttl int
}

//ttl in second
//expireCb,  expire callback
func NewTimerCache(ttl int, expireCb func(key, value interface{})) *TimerCache {
	tc := &TimerCache{
		q:   []*pqItem{},
		m:   map[interface{}]*pqItem{},
		ttl: ttl,
	}

	go func() {
		for {
			tc.tryPop(time.Now().Unix(), expireCb)
			time.Sleep(time.Second)
		}
	}()
	return tc
}

func (p *TimerCache) Put(key, val interface{}) bool {
	p.Lock()
	defer p.Unlock()

	ttl := time.Now().Unix() + int64(p.ttl)
	if old := p.m[key]; old == nil {
		item := &pqItem{
			key:   key,
			value: val,
			ttl:   ttl,
		}
		heap.Push(&p.q, item)
		p.m[key] = item
	} else {
		old.value = val
		old.ttl = ttl
		heap.Fix(&p.q, old.index)
	}
	return true
}

func (p *TimerCache) Get(key interface{}) interface{} {
	p.RLock()
	defer p.RUnlock()

	if item := p.m[key]; item == nil {
		return nil
	} else {
		return item.value
	}
}

func (p *TimerCache) Remove(key interface{}) interface{} {
	p.Lock()
	defer p.Unlock()

	if item, ok := p.m[key]; ok {
		item.dead = true //mark dead
		delete(p.m, key)
		return item.value
	} else {
		return nil
	}
}

func (p *TimerCache) Len() int {
	p.RLock()
	defer p.RUnlock()
	return len(p.m)
}

func (p *TimerCache) tryPop(tick int64, expireCb func(key, value interface{})) {
	p.Lock()
	defer p.Unlock()

	for p.q.Len() > 0 {
		item := p.q[0]
		if item.ttl > tick {
			//no expire items
			//			log.Println("no expire items", item.ttl, tick)
			return
		}

		item = heap.Pop(&p.q).(*pqItem)
		delete(p.m, item.key)

		//ignore dead item
		if !item.dead && expireCb != nil {
			go expireCb(item.key, item.value)
		}
	}
}

func (p *TimerCache) Loop(cb func(key, value interface{})) {
	if cb == nil {
		return
	}
	for _, item := range p.m {
		cb(item.key, item.value)
	}
}
<commit_msg>fix:thread safe<commit_after>package utee

import (
	"container/heap"
	"sync"
	"time"
)

// An Item is something we manage in a priority queue.
type pqItem struct {
	key   interface{}
	value interface{}
	ttl   int64 //unix timestamp , in second
	index int
	dead  bool //mark as dead,
}

// A priorityQueue implements heap.Interface and holds Items.
type priorityQueue []*pqItem

func (pq priorityQueue) Len() int { return len(pq) }

func (pq priorityQueue) Less(i, j int) bool {
	return pq[i].ttl < pq[j].ttl
}

func (pq priorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}

func (pq *priorityQueue) Push(x interface{}) {
	n := len(*pq)
	item := x.(*pqItem)
	item.index = n
	*pq = append(*pq, item)
}

func (pq *priorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	item.index = -1 // for safety
	*pq = old[0 : n-1]
	return item
}

type TimerCache struct {
	sync.RWMutex

	q   priorityQueue
	m   map[interface{}]*pqItem
	ttl int
}

//ttl in second
//expireCb,  expire callback
func NewTimerCache(ttl int, expireCb func(key, value interface{})) *TimerCache {
	tc := &TimerCache{
		q:   []*pqItem{},
		m:   map[interface{}]*pqItem{},
		ttl: ttl,
	}

	go func() {
		for {
			tc.tryPop(time.Now().Unix(), expireCb)
			time.Sleep(time.Second)
		}
	}()
	return tc
}

func (p *TimerCache) Put(key, val interface{}) bool {
	p.Lock()
	defer p.Unlock()

	ttl := time.Now().Unix() + int64(p.ttl)
	if old := p.m[key]; old == nil {
		item := &pqItem{
			key:   key,
			value: val,
			ttl:   ttl,
		}
		heap.Push(&p.q, item)
		p.m[key] = item
	} else {
		old.value = val
		old.ttl = ttl
		heap.Fix(&p.q, old.index)
	}
	return true
}

func (p *TimerCache) Get(key interface{}) interface{} {
	p.RLock()
	defer p.RUnlock()

	if item := p.m[key]; item == nil {
		return nil
	} else {
		return item.value
	}
}

func (p *TimerCache) Remove(key interface{}) interface{} {
	p.Lock()
	defer p.Unlock()

	if item, ok := p.m[key]; ok {
		item.dead = true //mark dead
		delete(p.m, key)
		return item.value
	} else {
		return nil
	}
}

func (p *TimerCache) Len() int {
	p.RLock()
	defer p.RUnlock()
	return len(p.m)
}

func (p *TimerCache) tryPop(tick int64, expireCb func(key, value interface{})) {
	p.Lock()
	defer p.Unlock()

	for p.q.Len() > 0 {
		item := p.q[0]
		if item.ttl > tick {
			//no expire items
			//			log.Println("no expire items", item.ttl, tick)
			return
		}

		item = heap.Pop(&p.q).(*pqItem)
		delete(p.m, item.key)

		//ignore dead item
		if !item.dead && expireCb != nil {
			go expireCb(item.key, item.value)
		}
	}
}

func (p *TimerCache) Loop(cb func(key, value interface{})) {
	if cb == nil {
		return
	}
	p.RLock()
	defer p.RUnlock()
	for _, item := range p.m {
		cb(item.key, item.value)
	}
}
<|endoftext|>


<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.

// Package token defines constants representing the lexical tokens of the Go
// programming language and basic operations on tokens (printing, predicates).
//
package token

import "strconv"

// Token is the set of lexical tokens of the Go programming language.
type Token int

// The list of tokens.
const (
	// Special tokens
	ILLEGAL Token = iota
	EOL
	EOF
	COMMENT

	literal_beg
	// Identifiers and basic type literals
	// (these tokens stand for classes of literals)
	IDENT  // main
	STRING // "abc"
	literal_end

	operator_beg
	// Operators and delimiters
	ASSIGN // =
	LBRACK // [
	RBRACK // ]
	operator_end
)

var tokens = [...]string{
	ILLEGAL: "ILLEGAL",

	EOL:     "EOL",
	EOF:     "EOF",
	COMMENT: "COMMENT",

	IDENT:  "IDENT",
	STRING: "STRING",

	ASSIGN: "=",
	LBRACK: "[",
	RBRACK: "]",
}

// String returns the string corresponding to the token tok.
// For operators, delimiters, and keywords the string is the actual
// token character sequence (e.g., for the token ADD, the string is
// "+"). For all other tokens the string corresponds to the token
// constant name (e.g. for the token IDENT, the string is "IDENT").
//
func (tok Token) String() string {
	s := ""
	if 0 <= tok && tok < Token(len(tokens)) {
		s = tokens[tok]
	}
	if s == "" {
		s = "token(" + strconv.Itoa(int(tok)) + ")"
	}
	return s
}

// A set of constants for precedence-based expression parsing.
// Non-operators have lowest precedence, followed by operators
// starting with precedence 1 up to unary operators. The highest
// precedence corresponds serves as "catch-all" precedence for
// selector, indexing, and other operator and delimiter tokens.
//
const (
	LowestPrec  = 0 // non-operators
	UnaryPrec   = 6
	HighestPrec = 7
)

// Precedence returns the operator precedence of the binary
// operator op. If op is not a binary operator, the result
// is LowestPrecedence.
//
func (op Token) Precedence() int {
	switch op {
	case LOR:
		return 1
	case LAND:
		return 2
	case EQL, NEQ, LSS, LEQ, GTR, GEQ:
		return 3
	case ADD, SUB, OR, XOR:
		return 4
	case MUL, QUO, REM, SHL, SHR, AND, AND_NOT:
		return 5
	}
	return LowestPrec
}

var keywords map[string]Token

func init() {
	keywords = make(map[string]Token)
	for i := keyword_beg + 1; i < keyword_end; i++ {
		keywords[tokens[i]] = i
	}
}

// Lookup maps an identifier to its keyword token or IDENT (if not a keyword).
//
func Lookup(ident string) Token {
	if tok, is_keyword := keywords[ident]; is_keyword {
		return tok
	}
	return IDENT
}

// Predicates

// IsLiteral returns true for tokens corresponding to identifiers
// and basic type literals; it returns false otherwise.
//
func (tok Token) IsLiteral() bool { return literal_beg < tok && tok < literal_end }

// IsOperator returns true for tokens corresponding to operators and
// delimiters; it returns false otherwise.
//
func (tok Token) IsOperator() bool { return operator_beg < tok && tok < operator_end }

// IsKeyword returns true for tokens corresponding to keywords;
// it returns false otherwise.
//
func (tok Token) IsKeyword() bool { return keyword_beg < tok && tok < keyword_end }
<commit_msg>token: remove Precedence, IsKeyword, Lookup<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.

// Package token defines constants representing the lexical tokens of the Go
// programming language and basic operations on tokens (printing, predicates).
//
package token

import "strconv"

// Token is the set of lexical tokens of the Go programming language.
type Token int

// The list of tokens.
const (
	// Special tokens
	ILLEGAL Token = iota
	EOL
	EOF
	COMMENT

	literal_beg
	// Identifiers and basic type literals
	// (these tokens stand for classes of literals)
	IDENT  // main
	STRING // "abc"
	literal_end

	operator_beg
	// Operators and delimiters
	ASSIGN // =
	LBRACK // [
	RBRACK // ]
	operator_end
)

var tokens = [...]string{
	ILLEGAL: "ILLEGAL",

	EOL:     "EOL",
	EOF:     "EOF",
	COMMENT: "COMMENT",

	IDENT:  "IDENT",
	STRING: "STRING",

	ASSIGN: "=",
	LBRACK: "[",
	RBRACK: "]",
}

// String returns the string corresponding to the token tok.
// For operators, delimiters, and keywords the string is the actual
// token character sequence (e.g., for the token ADD, the string is
// "+"). For all other tokens the string corresponds to the token
// constant name (e.g. for the token IDENT, the string is "IDENT").
//
func (tok Token) String() string {
	s := ""
	if 0 <= tok && tok < Token(len(tokens)) {
		s = tokens[tok]
	}
	if s == "" {
		s = "token(" + strconv.Itoa(int(tok)) + ")"
	}
	return s
}

// Predicates

// IsLiteral returns true for tokens corresponding to identifiers
// and basic type literals; it returns false otherwise.
//
func (tok Token) IsLiteral() bool { return literal_beg < tok && tok < literal_end }

// IsOperator returns true for tokens corresponding to operators and
// delimiters; it returns false otherwise.
//
func (tok Token) IsOperator() bool { return operator_beg < tok && tok < operator_end }
<|endoftext|>


<commit_before>package cq

import (
	"bytes"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const (
	transactionNotStarted = iota
	transactionStarted    = iota
)

type transactionResponse struct {
	Commit      string `json:"commit"`
	Transaction struct {
		Expires string
	}
	Errors []commitError `json:"errors"`
}

type cypherTransactionStatement struct {
	Statement  *string                `json:"statement"`
	Parameters map[string]interface{} `json:"parameters"`
}

type cypherTransaction struct {
	Statements     []cypherTransactionStatement `json:"statements"`
	commitURL      string
	transactionURL string
	expiration     time.Time
	c              *conn
	rows           []*rows
	keepAlive      *time.Timer
}

type commitError struct {
	Code    string `json:"code"`
	Message string `json:"message"`
}

type commitResponse struct {
	Errors []commitError `json:"errors"`
}

func (c *conn) Begin() (driver.Tx, error) {
	if c.transactionURL == "" {
		return nil, errTransactionsNotSupported
	}
	var buf bytes.Buffer
	json.NewEncoder(&buf).Encode(cypherTransaction{})
	req, err := http.NewRequest("POST", c.transactionURL, &buf)
	if err != nil {
		return nil, err
	}
	setDefaultHeaders(req)
	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	transResponse := transactionResponse{}
	json.NewDecoder(res.Body).Decode(&transResponse)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()
	if err != nil {
		return nil, err
	}
	exp, err := time.Parse(time.RFC1123Z, transResponse.Transaction.Expires)
	if err != nil {
		log.Println(err, c)
		err = nil
	}
	c.transaction = &cypherTransaction{
		commitURL:      transResponse.Commit,
		transactionURL: res.Header.Get("Location"),
		c:              c,
		expiration:     exp,
	}
	c.transaction.updateKeepAlive()
	return c.transaction, nil
}

func (tx *cypherTransaction) query(query *string, args []driver.Value) error {
	stmt := cypherTransactionStatement{
		Statement:  query,
		Parameters: makeArgsMap(args),
	}
	tx.Statements = append(tx.Statements, stmt)
	if len(tx.Statements) >= 100 {
		err := tx.exec()
		if err != nil {
			return err
		}
	}
	return nil
}

func (tx *cypherTransaction) exec() error {
	var buf bytes.Buffer
	err := json.NewEncoder(&buf).Encode(tx)
	if err != nil {
		return err
	}
	req, err := http.NewRequest("POST", tx.transactionURL, &buf)
	if err != nil {
		return err
	}
	setDefaultHeaders(req)
	res, err := client.Do(req)
	if err != nil {
		return err
	}
	trans := transactionResponse{}
	json.NewDecoder(res.Body).Decode(&trans)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()

	tx.expiration, err = time.Parse(time.RFC1123Z, trans.Transaction.Expires)
	if err != nil {
		log.Print(err, tx)
		err = nil
	}
	tx.updateKeepAlive()

	tx.Statements = tx.Statements[:0]

	if len(trans.Errors) > 0 {
		return errors.New("exec errors: " + fmt.Sprintf("%q", trans))
	}
	if err != nil {
		return err
	}
	return nil
}

func (tx *cypherTransaction) Commit() error {
	var buf bytes.Buffer
	err := json.NewEncoder(&buf).Encode(tx)
	if err != nil {
		return err
	}
	req, err := http.NewRequest("POST", tx.commitURL, &buf)
	if err != nil {
		return err
	}
	setDefaultHeaders(req)
	res, err := client.Do(req)
	defer res.Body.Close()
	if err != nil {
		return err
	}
	commit := commitResponse{}
	json.NewDecoder(res.Body).Decode(&commit)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()
	if err != nil {
		return err
	}
	if len(commit.Errors) > 0 {
		return errors.New("commit errors: " + fmt.Sprintf("%q", commit))
	}
	tx.c.transaction = nil
	if tx.keepAlive != nil {
		tx.keepAlive.Stop()
	}
	return nil
}

func (tx *cypherTransaction) Rollback() error {
	req, err := http.NewRequest("DELETE", tx.transactionURL, nil)
	if err != nil {
		return err
	}
	setDefaultHeaders(req)
	res, err := client.Do(req)
	if err != nil {
		return err
	}
	commit := commitResponse{}
	json.NewDecoder(res.Body).Decode(&commit)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()
	if err != nil {
		return err
	}
	if len(commit.Errors) > 0 {
		return errors.New("rollback errors: " + fmt.Sprintf("%q", commit))
	}
	tx.c.transaction = nil
	if tx.keepAlive != nil {
		tx.keepAlive.Stop()
	}

	return nil
}

func (tx *cypherTransaction) updateKeepAlive() {
	if tx.keepAlive != nil {
		tx.keepAlive.Stop()
	}
	dur := -1 * time.Since(tx.expiration)
	if dur <= 1*time.Second {
		dur = 500 * time.Millisecond
	}
	tx.keepAlive = time.AfterFunc(dur, func() { sendKeepAlive(tx.transactionURL) })
}

func sendKeepAlive(txURL string) {
	var buf bytes.Buffer
	err := json.NewEncoder(&buf).Encode(cypherTransaction{
		Statements: []cypherTransactionStatement{},
	})
	if err != nil {
		log.Print(err)
	}
	req, err := http.NewRequest("POST", txURL, &buf)
	if err != nil {
		log.Print(err)
	}
	setDefaultHeaders(req)
	res, err := client.Do(req)
	if err != nil {
	}
	trans := transactionResponse{}
	json.NewDecoder(res.Body).Decode(&trans)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()
	if len(trans.Errors) > 0 {
		//log.Println(trans)
		return
	}
	exp, err := time.Parse(time.RFC1123Z, trans.Transaction.Expires)
	if err != nil {
		//log.Println(err)
		return
	}
	dur := -1 * time.Since(exp)
	time.AfterFunc(dur, func() { sendKeepAlive(txURL) })
}
<commit_msg>collecting requests...<commit_after>package cq

import (
	"bytes"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"time"
)

const (
	transactionNotStarted = iota
	transactionStarted    = iota
)

type transactionResponse struct {
	Commit      string `json:"commit"`
	Transaction struct {
		Expires string
	}
	Errors   []commitError `json:"errors"`
	location string
}

type cypherTransactionStatement struct {
	Statement  *string                `json:"statement"`
	Parameters map[string]interface{} `json:"parameters"`
}

type cypherTransaction struct {
	Statements     []cypherTransactionStatement `json:"statements"`
	commitURL      string
	transactionURL string
	expiration     time.Time
	c              *conn
	rows           []*rows
	keepAlive      *time.Timer
}

type commitError struct {
	Code    string `json:"code"`
	Message string `json:"message"`
}

type commitResponse struct {
	Errors []commitError `json:"errors"`
}

func (c *conn) Begin() (driver.Tx, error) {
	if c.transactionURL == "" {
		return nil, errTransactionsNotSupported
	}
	transResponse, err := getTransactionResponse(c.transactionURL, cypherTransaction{})
	if err != nil {
		return nil, err
	}
	exp, err := time.Parse(time.RFC1123Z, transResponse.Transaction.Expires)
	if err != nil {
		log.Println(err, c)
		err = nil
	}
	c.transaction = &cypherTransaction{
		commitURL:      transResponse.Commit,
		transactionURL: transResponse.location,
		c:              c,
		expiration:     exp,
	}
	c.transaction.updateKeepAlive()
	return c.transaction, nil
}

func (tx *cypherTransaction) query(query *string, args []driver.Value) error {
	stmt := cypherTransactionStatement{
		Statement:  query,
		Parameters: makeArgsMap(args),
	}
	tx.Statements = append(tx.Statements, stmt)
	if len(tx.Statements) >= 100 {
		err := tx.exec()
		if err != nil {
			return err
		}
	}
	return nil
}

func (tx *cypherTransaction) exec() error {
	trans, err := getTransactionResponse(tx.transactionURL, *tx)
	if err != nil {
		return err
	}

	tx.expiration, err = time.Parse(time.RFC1123Z, trans.Transaction.Expires)
	if err != nil {
		log.Print(err, tx)
		err = nil
	}
	tx.updateKeepAlive()

	tx.Statements = tx.Statements[:0]

	if len(trans.Errors) > 0 {
		return errors.New("exec errors: " + fmt.Sprintf("%q", trans))
	}
	if err != nil {
		return err
	}
	return nil
}

func (tx *cypherTransaction) Commit() error {
	var buf bytes.Buffer
	err := json.NewEncoder(&buf).Encode(tx)
	if err != nil {
		return err
	}
	req, err := http.NewRequest("POST", tx.commitURL, &buf)
	if err != nil {
		return err
	}
	setDefaultHeaders(req)
	res, err := client.Do(req)
	defer res.Body.Close()
	if err != nil {
		return err
	}
	commit := commitResponse{}
	json.NewDecoder(res.Body).Decode(&commit)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()
	if err != nil {
		return err
	}
	if len(commit.Errors) > 0 {
		return errors.New("commit errors: " + fmt.Sprintf("%q", commit))
	}
	tx.c.transaction = nil
	if tx.keepAlive != nil {
		tx.keepAlive.Stop()
	}
	return nil
}

func (tx *cypherTransaction) Rollback() error {
	req, err := http.NewRequest("DELETE", tx.transactionURL, nil)
	if err != nil {
		return err
	}
	setDefaultHeaders(req)
	res, err := client.Do(req)
	if err != nil {
		return err
	}
	commit := commitResponse{}
	json.NewDecoder(res.Body).Decode(&commit)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()
	if err != nil {
		return err
	}
	if len(commit.Errors) > 0 {
		return errors.New("rollback errors: " + fmt.Sprintf("%q", commit))
	}
	tx.c.transaction = nil
	if tx.keepAlive != nil {
		tx.keepAlive.Stop()
	}

	return nil
}

func getTransactionResponse(url string, cypherTransReq cypherTransaction) (*transactionResponse, error) {
	var buf bytes.Buffer
	json.NewEncoder(&buf).Encode(cypherTransReq)
	req, err := http.NewRequest("POST", url, &buf)
	if err != nil {
		return nil, err
	}
	setDefaultHeaders(req)
	res, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	transResponse := transactionResponse{}
	json.NewDecoder(res.Body).Decode(&transResponse)
	io.Copy(ioutil.Discard, res.Body)
	res.Body.Close()
	transResponse.location = res.Header.Get("Location")
	return &transResponse, nil
}

func (tx *cypherTransaction) updateKeepAlive() {
	if tx.keepAlive != nil {
		tx.keepAlive.Stop()
	}
	dur := -1 * time.Since(tx.expiration)
	if dur <= 1*time.Second {
		dur = 300 * time.Millisecond
	}
	tx.keepAlive = time.AfterFunc(dur, func() { sendKeepAlive(tx.transactionURL) })
}

func sendKeepAlive(txURL string) {
	//log.Println("sending keepalive")
	trans, err := getTransactionResponse(txURL, cypherTransaction{Statements: []cypherTransactionStatement{}})
	if err != nil {
		//log.Println(trans)
		return
	}
	if len(trans.Errors) > 0 {
		//log.Println(trans)
		return
	}
	exp, err := time.Parse(time.RFC1123Z, trans.Transaction.Expires)
	if err != nil {
		//log.Println(err)
		return
	}
	dur := -1 * time.Since(exp)
	time.AfterFunc(dur, func() { sendKeepAlive(txURL) })
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 Dgraph Labs, Inc. and Contributors
 *
 * 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 badger

import (
	"bytes"
	"container/heap"
	"fmt"
	"strconv"
	"sync"
	"sync/atomic"

	"github.com/dgraph-io/badger/y"
	farm "github.com/dgryski/go-farm"
	"github.com/pkg/errors"
)

type uint64Heap []uint64

func (u uint64Heap) Len() int               { return len(u) }
func (u uint64Heap) Less(i int, j int) bool { return u[i] < u[j] }
func (u uint64Heap) Swap(i int, j int)      { u[i], u[j] = u[j], u[i] }
func (u *uint64Heap) Push(x interface{})    { *u = append(*u, x.(uint64)) }
func (u *uint64Heap) Pop() interface{} {
	old := *u
	n := len(old)
	x := old[n-1]
	*u = old[0 : n-1]
	return x
}

type globalTxnState struct {
	sync.Mutex
	curRead    uint64
	nextCommit uint64

	// These two structures are used to figure out when a commit is done. The minimum done commit is
	// used to update curRead.
	commitMark     uint64Heap
	pendingCommits map[uint64]struct{}

	// commits stores a key fingerprint and latest commit counter for it.
	// refCount is used to clear out commits map to avoid a memory blowup.
	commits  map[uint64]uint64
	refCount int64
}

func (gs *globalTxnState) addRef() {
	atomic.AddInt64(&gs.refCount, 1)
}

func (gs *globalTxnState) decrRef() {
	if count := atomic.AddInt64(&gs.refCount, -1); count == 0 {
		// Clear out pendingCommits maps to release memory.
		gs.Lock()
		y.AssertTrue(len(gs.commitMark) == 0)
		y.AssertTrue(len(gs.pendingCommits) == 0)
		if len(gs.commits) >= 1000 { // If the map is still small, let it slide.
			gs.commits = make(map[uint64]uint64)
		}
		gs.Unlock()
	}
}

func (gs *globalTxnState) readTs() uint64 {
	return atomic.LoadUint64(&gs.curRead)
}

func (gs *globalTxnState) commitTs() uint64 {
	gs.Lock()
	defer gs.Unlock()
	return gs.nextCommit
}

// hasConflict must be called while having a lock.
func (gs *globalTxnState) hasConflict(txn *Txn) bool {
	if len(txn.reads) == 0 {
		return false
	}
	for _, ro := range txn.reads {
		if ts, has := gs.commits[ro]; has && ts > txn.readTs {
			return true
		}
	}
	return false
}

func (gs *globalTxnState) newCommitTs(txn *Txn) uint64 {
	gs.Lock()
	defer gs.Unlock()

	if gs.hasConflict(txn) {
		return 0
	}

	var ts uint64
	if txn.commitTs == 0 {
		// This is the general case, when user doesn't specify the read and commit ts.
		ts = gs.nextCommit
		gs.nextCommit++

	} else {
		// If commitTs is set, use it instead.
		ts = txn.commitTs
		if gs.nextCommit <= ts { // Update this to max+1 commit ts, so replay works.
			gs.nextCommit = ts + 1
		}
	}

	for _, w := range txn.writes {
		gs.commits[w] = ts // Update the commitTs.
	}
	heap.Push(&gs.commitMark, ts)
	if _, has := gs.pendingCommits[ts]; has {
		panic(fmt.Sprintf("We shouldn't have the commit ts: %d", ts))
	}
	gs.pendingCommits[ts] = struct{}{}
	return ts
}

func (gs *globalTxnState) doneCommit(cts uint64) {
	gs.Lock()
	defer gs.Unlock()

	if _, has := gs.pendingCommits[cts]; !has {
		panic(fmt.Sprintf("We should already have the commit ts: %d", cts))
	}
	delete(gs.pendingCommits, cts)

	var min uint64
	for len(gs.commitMark) > 0 {
		ts := gs.commitMark[0]
		if _, has := gs.pendingCommits[ts]; has {
			// Still waiting for a txn to commit.
			break
		}
		min = ts
		heap.Pop(&gs.commitMark)
	}
	if min == 0 {
		return
	}
	atomic.StoreUint64(&gs.curRead, min)
	// nextCommit must never be reset.
}

type Txn struct {
	readTs   uint64
	commitTs uint64

	update bool     // update is used to conditionally keep track of reads.
	reads  []uint64 // contains fingerprints of keys read.
	writes []uint64 // contains fingerprints of keys written.

	pendingWrites map[string]*entry // cache stores any writes done by txn.

	kv        *KV
	callbacks []func()
	discarded bool
}

// Set sets the provided value for a given key. If key is not present, it is created.  If it is
// present, a new version is created at commit timestamp.
// Along with key and value, Set can also take an optional userMeta byte. This byte is stored
// alongside the key, and can be used as an aid to interpret the value or store other contextual
// bits corresponding to the key-value pair.
// This would fail with ErrReadOnlyTxn if update flag was set to false when creating this
// transaction.
func (txn *Txn) Set(key, val []byte, userMeta byte) error {
	if !txn.update {
		return ErrReadOnlyTxn
	} else if txn.discarded {
		return ErrDiscardedTxn
	} else if len(key) == 0 {
		return ErrEmptyKey
	} else if len(key) > maxKeySize {
		return exceedsMaxKeySizeError(key)
	} else if int64(len(val)) > txn.kv.opt.ValueLogFileSize {
		return exceedsMaxValueSizeError(val, txn.kv.opt.ValueLogFileSize)
	}

	fp := farm.Fingerprint64(key) // Avoid dealing with byte arrays.
	txn.writes = append(txn.writes, fp)

	e := &entry{
		Key:      key,
		Value:    val,
		UserMeta: userMeta,
	}
	txn.pendingWrites[string(key)] = e
	return nil
}

// Delete deletes a key. This is done by adding a delete marker for the key at commit timestamp.
// Any reads happening before this timestamp would be unaffected. Any reads after this commit would
// see the deletion.
func (txn *Txn) Delete(key []byte) error {
	if !txn.update {
		return ErrReadOnlyTxn
	} else if txn.discarded {
		return ErrDiscardedTxn
	} else if len(key) == 0 {
		return ErrEmptyKey
	} else if len(key) > maxKeySize {
		return exceedsMaxKeySizeError(key)
	}

	fp := farm.Fingerprint64(key) // Avoid dealing with byte arrays.
	txn.writes = append(txn.writes, fp)

	e := &entry{
		Key:  key,
		Meta: BitDelete,
	}
	txn.pendingWrites[string(key)] = e
	return nil
}

// Get looks for key and returns a KVItem.
// If key is not found, ErrKeyNotFound is returned.
func (txn *Txn) Get(key []byte) (item KVItem, rerr error) {
	if len(key) == 0 {
		return item, ErrEmptyKey
	} else if txn.discarded {
		return ErrDiscardedTxn
	}

	if txn.update {
		if e, has := txn.pendingWrites[string(key)]; has && bytes.Compare(key, e.Key) == 0 {
			// Fulfill from cache.
			item.meta = e.Meta
			item.val = e.Value
			item.userMeta = e.UserMeta
			item.key = key
			item.status = prefetched
			item.version = txn.readTs
			// We probably don't need to set KV on item here.
			return item, nil
		}
		// Only track reads if this is update txn. No need to track read if txn serviced it
		// internally.
		fp := farm.Fingerprint64(key)
		txn.reads = append(txn.reads, fp)
	}

	seek := y.KeyWithTs(key, txn.readTs)
	vs, err := txn.kv.get(seek)
	if err != nil {
		return item, errors.Wrapf(err, "KV::Get key: %q", key)
	}
	if vs.Value == nil && vs.Meta == 0 {
		return item, ErrKeyNotFound
	}
	if (vs.Meta & BitDelete) != 0 {
		return item, ErrKeyNotFound
	}

	item.key = key
	item.version = vs.Version
	item.meta = vs.Meta
	item.userMeta = vs.UserMeta
	item.kv = txn.kv
	item.vptr = vs.Value
	item.txn = txn
	return item, nil
}

func (txn *Txn) Discard() {
	if txn.discarded { // Avoid a re-run.
		return
	}
	txn.discarded = true

	for _, cb := range txn.callbacks {
		cb()
	}
	if txn.update {
		txn.kv.txnState.decrRef()
	}
}

// Commit commits the transaction, following these steps:
// 1. If there are no writes, return immediately.
// 2. Check if read rows were updated since txn started. If so, return ErrConflict.
// 3. If no conflict, generate a commit timestamp and update written rows' commit ts.
// 4. Batch up all writes, write them to value log and LSM tree.
// 5. If callback is provided, don't block on these writes, running this part asynchronously. The
// callback would be run once writes are done, along with any errors.
//
// If error is nil, the transaction is successfully committed. Else, Badger would automatically
// rollback the transaction, removing any writes from LSM tree. Note that the writes might be
// present on value log, but those can be garbage collected later.
func (txn *Txn) Commit(callback func(error)) error {
	if txn.discarded {
		return ErrDiscardedTxn
	}
	defer txn.Discard()
	if len(txn.writes) == 0 {
		return nil // Nothing to do.
	}

	state := txn.kv.txnState
	commitTs := state.newCommitTs(txn)
	if commitTs == 0 {
		return ErrConflict
	}
	defer state.doneCommit(commitTs)

	entries := make([]*entry, 0, len(txn.pendingWrites)+1)
	for _, e := range txn.pendingWrites {
		// Suffix the keys with commit ts, so the key versions are sorted in
		// descending order of commit timestamp.
		e.Key = y.KeyWithTs(e.Key, commitTs)
		e.Meta |= BitTxn
		entries = append(entries, e)
	}
	e := &entry{
		Key:   y.KeyWithTs(txnKey, commitTs),
		Value: []byte(strconv.FormatUint(commitTs, 10)),
		Meta:  BitFinTxn,
	}
	entries = append(entries, e)

	if callback == nil {
		// If batchSet failed, LSM would not have been updated. So, no need to rollback anything.

		// TODO: What if some of the txns successfully make it to value log, but others fail.
		// Nothing gets updated to LSM, until a restart happens.
		return txn.kv.batchSet(entries)
	}
	return txn.kv.batchSetAsync(entries, callback)
}

func (txn *Txn) CommitAt(commitTs uint64, callback func(error)) error {
	txn.commitTs = commitTs
	return txn.Commit(callback)
}

// NewIterator returns a new iterator. Depending upon the options, either only keys, or both
// key-value pairs would be fetched. The keys are returned in lexicographically sorted order.
// Usage:
//   opt := badger.DefaultIteratorOptions
//   itr := txn.NewIterator(opt)
//   for itr.Rewind(); itr.Valid(); itr.Next() {
//     item := itr.Item()
//     key := item.Key()
//     var val []byte
//     err = item.Value(func(v []byte) {
//         val = make([]byte, len(v))
// 	       copy(val, v)
//     }) 	// This could block while value is fetched from value log.
//          // For key only iteration, set opt.PrefetchValues to false, and don't call
//          // item.Value(func(v []byte)).
//
//     // Remember that both key, val would become invalid in the next iteration of the loop.
//     // So, if you need access to them outside, copy them or parse them.
//   }
//   itr.Close()
// TODO: Move this usage to README.
func (txn *Txn) NewIterator(opt IteratorOptions) *Iterator {
	tables, decr := txn.kv.getMemTables()
	defer decr()
	txn.kv.vlog.incrIteratorCount()
	var iters []y.Iterator
	for i := 0; i < len(tables); i++ {
		iters = append(iters, tables[i].NewUniIterator(opt.Reverse))
	}
	iters = txn.kv.lc.appendIterators(iters, opt.Reverse) // This will increment references.
	res := &Iterator{
		txn:    txn,
		iitr:   y.NewMergeIterator(iters, opt.Reverse),
		opt:    opt,
		readTs: txn.readTs,
	}
	return res
}

func (kv *KV) NewTransaction(update bool) *Txn {
	txn := &Txn{
		update: update,
		kv:     kv,
		readTs: kv.txnState.readTs(),
	}
	if update {
		txn.pendingWrites = make(map[string]*entry)
		txn.kv.txnState.addRef()
	}
	return txn
}

func (kv *KV) NewTransactionAt(readTs uint64, update bool) *Txn {
	txn := kv.NewTransaction(update)
	txn.readTs = readTs
	return txn
}

func (kv *KV) View(callback func(txn *Txn) error) error {
	txn := kv.NewTransaction(false)
	defer txn.Discard()

	return callback(txn)
}
<commit_msg>Fix return in Get<commit_after>/*
 * Copyright 2017 Dgraph Labs, Inc. and Contributors
 *
 * 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 badger

import (
	"bytes"
	"container/heap"
	"fmt"
	"strconv"
	"sync"
	"sync/atomic"

	"github.com/dgraph-io/badger/y"
	farm "github.com/dgryski/go-farm"
	"github.com/pkg/errors"
)

type uint64Heap []uint64

func (u uint64Heap) Len() int               { return len(u) }
func (u uint64Heap) Less(i int, j int) bool { return u[i] < u[j] }
func (u uint64Heap) Swap(i int, j int)      { u[i], u[j] = u[j], u[i] }
func (u *uint64Heap) Push(x interface{})    { *u = append(*u, x.(uint64)) }
func (u *uint64Heap) Pop() interface{} {
	old := *u
	n := len(old)
	x := old[n-1]
	*u = old[0 : n-1]
	return x
}

type globalTxnState struct {
	sync.Mutex
	curRead    uint64
	nextCommit uint64

	// These two structures are used to figure out when a commit is done. The minimum done commit is
	// used to update curRead.
	commitMark     uint64Heap
	pendingCommits map[uint64]struct{}

	// commits stores a key fingerprint and latest commit counter for it.
	// refCount is used to clear out commits map to avoid a memory blowup.
	commits  map[uint64]uint64
	refCount int64
}

func (gs *globalTxnState) addRef() {
	atomic.AddInt64(&gs.refCount, 1)
}

func (gs *globalTxnState) decrRef() {
	if count := atomic.AddInt64(&gs.refCount, -1); count == 0 {
		// Clear out pendingCommits maps to release memory.
		gs.Lock()
		y.AssertTrue(len(gs.commitMark) == 0)
		y.AssertTrue(len(gs.pendingCommits) == 0)
		if len(gs.commits) >= 1000 { // If the map is still small, let it slide.
			gs.commits = make(map[uint64]uint64)
		}
		gs.Unlock()
	}
}

func (gs *globalTxnState) readTs() uint64 {
	return atomic.LoadUint64(&gs.curRead)
}

func (gs *globalTxnState) commitTs() uint64 {
	gs.Lock()
	defer gs.Unlock()
	return gs.nextCommit
}

// hasConflict must be called while having a lock.
func (gs *globalTxnState) hasConflict(txn *Txn) bool {
	if len(txn.reads) == 0 {
		return false
	}
	for _, ro := range txn.reads {
		if ts, has := gs.commits[ro]; has && ts > txn.readTs {
			return true
		}
	}
	return false
}

func (gs *globalTxnState) newCommitTs(txn *Txn) uint64 {
	gs.Lock()
	defer gs.Unlock()

	if gs.hasConflict(txn) {
		return 0
	}

	var ts uint64
	if txn.commitTs == 0 {
		// This is the general case, when user doesn't specify the read and commit ts.
		ts = gs.nextCommit
		gs.nextCommit++

	} else {
		// If commitTs is set, use it instead.
		ts = txn.commitTs
		if gs.nextCommit <= ts { // Update this to max+1 commit ts, so replay works.
			gs.nextCommit = ts + 1
		}
	}

	for _, w := range txn.writes {
		gs.commits[w] = ts // Update the commitTs.
	}
	heap.Push(&gs.commitMark, ts)
	if _, has := gs.pendingCommits[ts]; has {
		panic(fmt.Sprintf("We shouldn't have the commit ts: %d", ts))
	}
	gs.pendingCommits[ts] = struct{}{}
	return ts
}

func (gs *globalTxnState) doneCommit(cts uint64) {
	gs.Lock()
	defer gs.Unlock()

	if _, has := gs.pendingCommits[cts]; !has {
		panic(fmt.Sprintf("We should already have the commit ts: %d", cts))
	}
	delete(gs.pendingCommits, cts)

	var min uint64
	for len(gs.commitMark) > 0 {
		ts := gs.commitMark[0]
		if _, has := gs.pendingCommits[ts]; has {
			// Still waiting for a txn to commit.
			break
		}
		min = ts
		heap.Pop(&gs.commitMark)
	}
	if min == 0 {
		return
	}
	atomic.StoreUint64(&gs.curRead, min)
	// nextCommit must never be reset.
}

type Txn struct {
	readTs   uint64
	commitTs uint64

	update bool     // update is used to conditionally keep track of reads.
	reads  []uint64 // contains fingerprints of keys read.
	writes []uint64 // contains fingerprints of keys written.

	pendingWrites map[string]*entry // cache stores any writes done by txn.

	kv        *KV
	callbacks []func()
	discarded bool
}

// Set sets the provided value for a given key. If key is not present, it is created.  If it is
// present, a new version is created at commit timestamp.
// Along with key and value, Set can also take an optional userMeta byte. This byte is stored
// alongside the key, and can be used as an aid to interpret the value or store other contextual
// bits corresponding to the key-value pair.
// This would fail with ErrReadOnlyTxn if update flag was set to false when creating this
// transaction.
func (txn *Txn) Set(key, val []byte, userMeta byte) error {
	if !txn.update {
		return ErrReadOnlyTxn
	} else if txn.discarded {
		return ErrDiscardedTxn
	} else if len(key) == 0 {
		return ErrEmptyKey
	} else if len(key) > maxKeySize {
		return exceedsMaxKeySizeError(key)
	} else if int64(len(val)) > txn.kv.opt.ValueLogFileSize {
		return exceedsMaxValueSizeError(val, txn.kv.opt.ValueLogFileSize)
	}

	fp := farm.Fingerprint64(key) // Avoid dealing with byte arrays.
	txn.writes = append(txn.writes, fp)

	e := &entry{
		Key:      key,
		Value:    val,
		UserMeta: userMeta,
	}
	txn.pendingWrites[string(key)] = e
	return nil
}

// Delete deletes a key. This is done by adding a delete marker for the key at commit timestamp.
// Any reads happening before this timestamp would be unaffected. Any reads after this commit would
// see the deletion.
func (txn *Txn) Delete(key []byte) error {
	if !txn.update {
		return ErrReadOnlyTxn
	} else if txn.discarded {
		return ErrDiscardedTxn
	} else if len(key) == 0 {
		return ErrEmptyKey
	} else if len(key) > maxKeySize {
		return exceedsMaxKeySizeError(key)
	}

	fp := farm.Fingerprint64(key) // Avoid dealing with byte arrays.
	txn.writes = append(txn.writes, fp)

	e := &entry{
		Key:  key,
		Meta: BitDelete,
	}
	txn.pendingWrites[string(key)] = e
	return nil
}

// Get looks for key and returns a KVItem.
// If key is not found, ErrKeyNotFound is returned.
func (txn *Txn) Get(key []byte) (item KVItem, rerr error) {
	if len(key) == 0 {
		return item, ErrEmptyKey
	} else if txn.discarded {
		return item, ErrDiscardedTxn
	}

	if txn.update {
		if e, has := txn.pendingWrites[string(key)]; has && bytes.Compare(key, e.Key) == 0 {
			// Fulfill from cache.
			item.meta = e.Meta
			item.val = e.Value
			item.userMeta = e.UserMeta
			item.key = key
			item.status = prefetched
			item.version = txn.readTs
			// We probably don't need to set KV on item here.
			return item, nil
		}
		// Only track reads if this is update txn. No need to track read if txn serviced it
		// internally.
		fp := farm.Fingerprint64(key)
		txn.reads = append(txn.reads, fp)
	}

	seek := y.KeyWithTs(key, txn.readTs)
	vs, err := txn.kv.get(seek)
	if err != nil {
		return item, errors.Wrapf(err, "KV::Get key: %q", key)
	}
	if vs.Value == nil && vs.Meta == 0 {
		return item, ErrKeyNotFound
	}
	if (vs.Meta & BitDelete) != 0 {
		return item, ErrKeyNotFound
	}

	item.key = key
	item.version = vs.Version
	item.meta = vs.Meta
	item.userMeta = vs.UserMeta
	item.kv = txn.kv
	item.vptr = vs.Value
	item.txn = txn
	return item, nil
}

func (txn *Txn) Discard() {
	if txn.discarded { // Avoid a re-run.
		return
	}
	txn.discarded = true

	for _, cb := range txn.callbacks {
		cb()
	}
	if txn.update {
		txn.kv.txnState.decrRef()
	}
}

// Commit commits the transaction, following these steps:
// 1. If there are no writes, return immediately.
// 2. Check if read rows were updated since txn started. If so, return ErrConflict.
// 3. If no conflict, generate a commit timestamp and update written rows' commit ts.
// 4. Batch up all writes, write them to value log and LSM tree.
// 5. If callback is provided, don't block on these writes, running this part asynchronously. The
// callback would be run once writes are done, along with any errors.
//
// If error is nil, the transaction is successfully committed. Else, Badger would automatically
// rollback the transaction, removing any writes from LSM tree. Note that the writes might be
// present on value log, but those can be garbage collected later.
func (txn *Txn) Commit(callback func(error)) error {
	if txn.discarded {
		return ErrDiscardedTxn
	}
	defer txn.Discard()
	if len(txn.writes) == 0 {
		return nil // Nothing to do.
	}

	state := txn.kv.txnState
	commitTs := state.newCommitTs(txn)
	if commitTs == 0 {
		return ErrConflict
	}
	defer state.doneCommit(commitTs)

	entries := make([]*entry, 0, len(txn.pendingWrites)+1)
	for _, e := range txn.pendingWrites {
		// Suffix the keys with commit ts, so the key versions are sorted in
		// descending order of commit timestamp.
		e.Key = y.KeyWithTs(e.Key, commitTs)
		e.Meta |= BitTxn
		entries = append(entries, e)
	}
	e := &entry{
		Key:   y.KeyWithTs(txnKey, commitTs),
		Value: []byte(strconv.FormatUint(commitTs, 10)),
		Meta:  BitFinTxn,
	}
	entries = append(entries, e)

	if callback == nil {
		// If batchSet failed, LSM would not have been updated. So, no need to rollback anything.

		// TODO: What if some of the txns successfully make it to value log, but others fail.
		// Nothing gets updated to LSM, until a restart happens.
		return txn.kv.batchSet(entries)
	}
	return txn.kv.batchSetAsync(entries, callback)
}

func (txn *Txn) CommitAt(commitTs uint64, callback func(error)) error {
	txn.commitTs = commitTs
	return txn.Commit(callback)
}

// NewIterator returns a new iterator. Depending upon the options, either only keys, or both
// key-value pairs would be fetched. The keys are returned in lexicographically sorted order.
// Usage:
//   opt := badger.DefaultIteratorOptions
//   itr := txn.NewIterator(opt)
//   for itr.Rewind(); itr.Valid(); itr.Next() {
//     item := itr.Item()
//     key := item.Key()
//     var val []byte
//     err = item.Value(func(v []byte) {
//         val = make([]byte, len(v))
// 	       copy(val, v)
//     }) 	// This could block while value is fetched from value log.
//          // For key only iteration, set opt.PrefetchValues to false, and don't call
//          // item.Value(func(v []byte)).
//
//     // Remember that both key, val would become invalid in the next iteration of the loop.
//     // So, if you need access to them outside, copy them or parse them.
//   }
//   itr.Close()
// TODO: Move this usage to README.
func (txn *Txn) NewIterator(opt IteratorOptions) *Iterator {
	tables, decr := txn.kv.getMemTables()
	defer decr()
	txn.kv.vlog.incrIteratorCount()
	var iters []y.Iterator
	for i := 0; i < len(tables); i++ {
		iters = append(iters, tables[i].NewUniIterator(opt.Reverse))
	}
	iters = txn.kv.lc.appendIterators(iters, opt.Reverse) // This will increment references.
	res := &Iterator{
		txn:    txn,
		iitr:   y.NewMergeIterator(iters, opt.Reverse),
		opt:    opt,
		readTs: txn.readTs,
	}
	return res
}

func (kv *KV) NewTransaction(update bool) *Txn {
	txn := &Txn{
		update: update,
		kv:     kv,
		readTs: kv.txnState.readTs(),
	}
	if update {
		txn.pendingWrites = make(map[string]*entry)
		txn.kv.txnState.addRef()
	}
	return txn
}

func (kv *KV) NewTransactionAt(readTs uint64, update bool) *Txn {
	txn := kv.NewTransaction(update)
	txn.readTs = readTs
	return txn
}

func (kv *KV) View(callback func(txn *Txn) error) error {
	txn := kv.NewTransaction(false)
	defer txn.Discard()

	return callback(txn)
}
<|endoftext|>