<commit_before>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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

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

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

package scheduler

import (
	"errors"
	"math/rand"
	"reflect"
	"testing"
	"time"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/testapi"
	"k8s.io/kubernetes/pkg/client/cache"
	"k8s.io/kubernetes/pkg/client/record"
	"k8s.io/kubernetes/pkg/util"
	"k8s.io/kubernetes/plugin/pkg/scheduler/algorithm"
	"k8s.io/kubernetes/plugin/pkg/scheduler/algorithm/predicates"
)

type fakeBinder struct {
	b func(binding *api.Binding) error
}

func (fb fakeBinder) Bind(binding *api.Binding) error { return fb.b(binding) }

func podWithID(id, desiredHost string) *api.Pod {
	return &api.Pod{
		ObjectMeta: api.ObjectMeta{Name: id, SelfLink: testapi.Default.SelfLink("pods", id)},
		Spec: api.PodSpec{
			NodeName: desiredHost,
		},
	}
}

func podWithPort(id, desiredHost string, port int) *api.Pod {
	pod := podWithID(id, desiredHost)
	pod.Spec.Containers = []api.Container{
		{Name: "ctr", Ports: []api.ContainerPort{{HostPort: port}}},
	}
	return pod
}

type mockScheduler struct {
	machine string
	err     error
}

func (es mockScheduler) Schedule(pod *api.Pod, ml algorithm.NodeLister) (string, error) {
	return es.machine, es.err
}

func TestScheduler(t *testing.T) {
	eventBroadcaster := record.NewBroadcaster()
	defer eventBroadcaster.StartLogging(t.Logf).Stop()
	errS := errors.New("scheduler")
	errB := errors.New("binder")

	table := []struct {
		injectBindError  error
		sendPod          *api.Pod
		algo             algorithm.ScheduleAlgorithm
		expectErrorPod   *api.Pod
		expectAssumedPod *api.Pod
		expectError      error
		expectBind       *api.Binding
		eventReason      string
	}{
		{
			sendPod:          podWithID("foo", ""),
			algo:             mockScheduler{"machine1", nil},
			expectBind:       &api.Binding{ObjectMeta: api.ObjectMeta{Name: "foo"}, Target: api.ObjectReference{Kind: "Node", Name: "machine1"}},
			expectAssumedPod: podWithID("foo", "machine1"),
			eventReason:      "Scheduled",
		}, {
			sendPod:        podWithID("foo", ""),
			algo:           mockScheduler{"machine1", errS},
			expectError:    errS,
			expectErrorPod: podWithID("foo", ""),
			eventReason:    "FailedScheduling",
		}, {
			sendPod:         podWithID("foo", ""),
			algo:            mockScheduler{"machine1", nil},
			expectBind:      &api.Binding{ObjectMeta: api.ObjectMeta{Name: "foo"}, Target: api.ObjectReference{Kind: "Node", Name: "machine1"}},
			injectBindError: errB,
			expectError:     errB,
			expectErrorPod:  podWithID("foo", ""),
			eventReason:     "FailedScheduling",
		},
	}

	for i, item := range table {
		var gotError error
		var gotPod *api.Pod
		var gotAssumedPod *api.Pod
		var gotBinding *api.Binding
		c := &Config{
			Modeler: &FakeModeler{
				AssumePodFunc: func(pod *api.Pod) {
					gotAssumedPod = pod
				},
			},
			NodeLister: algorithm.FakeNodeLister(
				api.NodeList{Items: []api.Node{{ObjectMeta: api.ObjectMeta{Name: "machine1"}}}},
			),
			Algorithm: item.algo,
			Binder: fakeBinder{func(b *api.Binding) error {
				gotBinding = b
				return item.injectBindError
			}},
			Error: func(p *api.Pod, err error) {
				gotPod = p
				gotError = err
			},
			NextPod: func() *api.Pod {
				return item.sendPod
			},
			Recorder: eventBroadcaster.NewRecorder(api.EventSource{Component: "scheduler"}),
		}
		s := New(c)
		called := make(chan struct{})
		events := eventBroadcaster.StartEventWatcher(func(e *api.Event) {
			if e, a := item.eventReason, e.Reason; e != a {
				t.Errorf("%v: expected %v, got %v", i, e, a)
			}
			close(called)
		})
		s.scheduleOne()
		if e, a := item.expectAssumedPod, gotAssumedPod; !reflect.DeepEqual(e, a) {
			t.Errorf("%v: assumed pod: wanted %v, got %v", i, e, a)
		}
		if e, a := item.expectErrorPod, gotPod; !reflect.DeepEqual(e, a) {
			t.Errorf("%v: error pod: wanted %v, got %v", i, e, a)
		}
		if e, a := item.expectError, gotError; !reflect.DeepEqual(e, a) {
			t.Errorf("%v: error: wanted %v, got %v", i, e, a)
		}
		if e, a := item.expectBind, gotBinding; !reflect.DeepEqual(e, a) {
			t.Errorf("%v: error: %s", i, util.ObjectDiff(e, a))
		}
		<-called
		events.Stop()
	}
}

func TestSchedulerForgetAssumedPodAfterDelete(t *testing.T) {
	eventBroadcaster := record.NewBroadcaster()
	defer eventBroadcaster.StartLogging(t.Logf).Stop()

	// Setup modeler so we control the contents of all 3 stores: assumed,
	// scheduled and queued
	scheduledPodStore := cache.NewStore(cache.MetaNamespaceKeyFunc)
	scheduledPodLister := &cache.StoreToPodLister{Store: scheduledPodStore}

	queuedPodStore := cache.NewFIFO(cache.MetaNamespaceKeyFunc)
	queuedPodLister := &cache.StoreToPodLister{Store: queuedPodStore}

	modeler := NewSimpleModeler(queuedPodLister, scheduledPodLister)

	// Create a fake clock used to timestamp entries and calculate ttl. Nothing
	// will expire till we flip to something older than the ttl, at which point
	// all entries inserted with fakeTime will expire.
	ttl := 30 * time.Second
	fakeTime := time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC)
	fakeClock := util.NewFakeClock(fakeTime)
	ttlPolicy := &cache.TTLPolicy{Ttl: ttl, Clock: fakeClock}
	assumedPodsStore := cache.NewFakeExpirationStore(
		cache.MetaNamespaceKeyFunc, nil, ttlPolicy, fakeClock)
	modeler.assumedPods = &cache.StoreToPodLister{Store: assumedPodsStore}

	// Port is the easiest way to cause a fit predicate failure
	podPort := 8080
	firstPod := podWithPort("foo", "", podPort)

	// Create the scheduler config
	algo := NewGenericScheduler(
		map[string]algorithm.FitPredicate{"PodFitsHostPorts": predicates.PodFitsHostPorts},
		[]algorithm.PriorityConfig{},
		[]algorithm.SchedulerExtender{},
		modeler.PodLister(),
		rand.New(rand.NewSource(time.Now().UnixNano())))

	var gotBinding *api.Binding
	c := &Config{
		Modeler: modeler,
		NodeLister: algorithm.FakeNodeLister(
			api.NodeList{Items: []api.Node{{ObjectMeta: api.ObjectMeta{Name: "machine1"}}}},
		),
		Algorithm: algo,
		Binder: fakeBinder{func(b *api.Binding) error {
			scheduledPodStore.Add(podWithPort(b.Name, b.Target.Name, podPort))
			gotBinding = b
			return nil
		}},
		NextPod: func() *api.Pod {
			return queuedPodStore.Pop().(*api.Pod)
		},
		Error: func(p *api.Pod, err error) {
			t.Errorf("Unexpected error when scheduling pod %+v: %v", p, err)
		},
		Recorder: eventBroadcaster.NewRecorder(api.EventSource{Component: "scheduler"}),
	}

	// First scheduling pass should schedule the pod
	s := New(c)
	called := make(chan struct{})
	events := eventBroadcaster.StartEventWatcher(func(e *api.Event) {
		if e, a := "Scheduled", e.Reason; e != a {
			t.Errorf("expected %v, got %v", e, a)
		}
		close(called)
	})

	queuedPodStore.Add(firstPod)
	// queuedPodStore: [foo:8080]
	// scheduledPodStore: []
	// assumedPods: []

	s.scheduleOne()
	// queuedPodStore: []
	// scheduledPodStore: [foo:8080]
	// assumedPods: [foo:8080]

	pod, exists, _ := scheduledPodStore.GetByKey("foo")
	if !exists {
		t.Errorf("Expected scheduled pod store to contain pod")
	}
	pod, exists, _ = queuedPodStore.GetByKey("foo")
	if exists {
		t.Errorf("Did not expect a queued pod, found %+v", pod)
	}
	pod, exists, _ = assumedPodsStore.GetByKey("foo")
	if !exists {
		t.Errorf("Assumed pod store should contain stale pod")
	}

	expectBind := &api.Binding{
		ObjectMeta: api.ObjectMeta{Name: "foo"},
		Target:     api.ObjectReference{Kind: "Node", Name: "machine1"},
	}
	if ex, ac := expectBind, gotBinding; !reflect.DeepEqual(ex, ac) {
		t.Errorf("Expected exact match on binding: %s", util.ObjectDiff(ex, ac))
	}

	<-called
	events.Stop()

	scheduledPodStore.Delete(pod)
	_, exists, _ = assumedPodsStore.Get(pod)
	if !exists {
		t.Errorf("Expected pod %#v in assumed pod store", pod)
	}

	secondPod := podWithPort("bar", "", podPort)
	queuedPodStore.Add(secondPod)
	// queuedPodStore: [bar:8080]
	// scheduledPodStore: []
	// assumedPods: [foo:8080]

	// Second scheduling pass will fail to schedule if the store hasn't expired
	// the deleted pod. This would normally happen with a timeout.
	//expirationPolicy.NeverExpire = util.NewStringSet()
	fakeClock.Step(ttl + 1)

	called = make(chan struct{})
	events = eventBroadcaster.StartEventWatcher(func(e *api.Event) {
		if e, a := "Scheduled", e.Reason; e != a {
			t.Errorf("expected %v, got %v", e, a)
		}
		close(called)
	})

	s.scheduleOne()

	expectBind = &api.Binding{
		ObjectMeta: api.ObjectMeta{Name: "bar"},
		Target:     api.ObjectReference{Kind: "Node", Name: "machine1"},
	}
	if ex, ac := expectBind, gotBinding; !reflect.DeepEqual(ex, ac) {
		t.Errorf("Expected exact match on binding: %s", util.ObjectDiff(ex, ac))
	}
	<-called
	events.Stop()
}
<commit_msg>Guarantee we finish writing to the test log before the test exits.<commit_after>/*
Copyright 2014 The Kubernetes Authors All rights reserved.

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

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

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

package scheduler

import (
	"errors"
	"fmt"
	"math/rand"
	"reflect"
	"sync"
	"testing"
	"time"

	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/api/testapi"
	"k8s.io/kubernetes/pkg/client/cache"
	"k8s.io/kubernetes/pkg/client/record"
	"k8s.io/kubernetes/pkg/util"
	"k8s.io/kubernetes/plugin/pkg/scheduler/algorithm"
	"k8s.io/kubernetes/plugin/pkg/scheduler/algorithm/predicates"
)

type fakeBinder struct {
	b func(binding *api.Binding) error
}

func (fb fakeBinder) Bind(binding *api.Binding) error { return fb.b(binding) }

func podWithID(id, desiredHost string) *api.Pod {
	return &api.Pod{
		ObjectMeta: api.ObjectMeta{Name: id, SelfLink: testapi.Default.SelfLink("pods", id)},
		Spec: api.PodSpec{
			NodeName: desiredHost,
		},
	}
}

func podWithPort(id, desiredHost string, port int) *api.Pod {
	pod := podWithID(id, desiredHost)
	pod.Spec.Containers = []api.Container{
		{Name: "ctr", Ports: []api.ContainerPort{{HostPort: port}}},
	}
	return pod
}

type mockScheduler struct {
	machine string
	err     error
}

func (es mockScheduler) Schedule(pod *api.Pod, ml algorithm.NodeLister) (string, error) {
	return es.machine, es.err
}

func TestScheduler(t *testing.T) {
	eventBroadcaster := record.NewBroadcaster()
	defer eventBroadcaster.StartLogging(t.Logf).Stop()
	errS := errors.New("scheduler")
	errB := errors.New("binder")

	table := []struct {
		injectBindError  error
		sendPod          *api.Pod
		algo             algorithm.ScheduleAlgorithm
		expectErrorPod   *api.Pod
		expectAssumedPod *api.Pod
		expectError      error
		expectBind       *api.Binding
		eventReason      string
	}{
		{
			sendPod:          podWithID("foo", ""),
			algo:             mockScheduler{"machine1", nil},
			expectBind:       &api.Binding{ObjectMeta: api.ObjectMeta{Name: "foo"}, Target: api.ObjectReference{Kind: "Node", Name: "machine1"}},
			expectAssumedPod: podWithID("foo", "machine1"),
			eventReason:      "Scheduled",
		}, {
			sendPod:        podWithID("foo", ""),
			algo:           mockScheduler{"machine1", errS},
			expectError:    errS,
			expectErrorPod: podWithID("foo", ""),
			eventReason:    "FailedScheduling",
		}, {
			sendPod:         podWithID("foo", ""),
			algo:            mockScheduler{"machine1", nil},
			expectBind:      &api.Binding{ObjectMeta: api.ObjectMeta{Name: "foo"}, Target: api.ObjectReference{Kind: "Node", Name: "machine1"}},
			injectBindError: errB,
			expectError:     errB,
			expectErrorPod:  podWithID("foo", ""),
			eventReason:     "FailedScheduling",
		},
	}

	for i, item := range table {
		var gotError error
		var gotPod *api.Pod
		var gotAssumedPod *api.Pod
		var gotBinding *api.Binding
		c := &Config{
			Modeler: &FakeModeler{
				AssumePodFunc: func(pod *api.Pod) {
					gotAssumedPod = pod
				},
			},
			NodeLister: algorithm.FakeNodeLister(
				api.NodeList{Items: []api.Node{{ObjectMeta: api.ObjectMeta{Name: "machine1"}}}},
			),
			Algorithm: item.algo,
			Binder: fakeBinder{func(b *api.Binding) error {
				gotBinding = b
				return item.injectBindError
			}},
			Error: func(p *api.Pod, err error) {
				gotPod = p
				gotError = err
			},
			NextPod: func() *api.Pod {
				return item.sendPod
			},
			Recorder: eventBroadcaster.NewRecorder(api.EventSource{Component: "scheduler"}),
		}
		s := New(c)
		called := make(chan struct{})
		events := eventBroadcaster.StartEventWatcher(func(e *api.Event) {
			if e, a := item.eventReason, e.Reason; e != a {
				t.Errorf("%v: expected %v, got %v", i, e, a)
			}
			close(called)
		})
		s.scheduleOne()
		if e, a := item.expectAssumedPod, gotAssumedPod; !reflect.DeepEqual(e, a) {
			t.Errorf("%v: assumed pod: wanted %v, got %v", i, e, a)
		}
		if e, a := item.expectErrorPod, gotPod; !reflect.DeepEqual(e, a) {
			t.Errorf("%v: error pod: wanted %v, got %v", i, e, a)
		}
		if e, a := item.expectError, gotError; !reflect.DeepEqual(e, a) {
			t.Errorf("%v: error: wanted %v, got %v", i, e, a)
		}
		if e, a := item.expectBind, gotBinding; !reflect.DeepEqual(e, a) {
			t.Errorf("%v: error: %s", i, util.ObjectDiff(e, a))
		}
		<-called
		events.Stop()
	}
}

func TestSchedulerForgetAssumedPodAfterDelete(t *testing.T) {
	// Set up a channel through which we'll funnel log messages from the watcher.
	// This way, we can guarantee that when the test ends no thread will still be
	// trying to write to t.Logf (which it would if we handed t.Logf directly to
	// StartLogging).
	ch := make(chan string)
	done := make(chan struct{})
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		for {
			select {
			case msg := <-ch:
				t.Log(msg)
			case <-done:
				return
			}
		}
	}()
	eventBroadcaster := record.NewBroadcaster()
	watcher := eventBroadcaster.StartLogging(func(format string, args ...interface{}) {
		ch <- fmt.Sprintf(format, args...)
	})
	defer func() {
		watcher.Stop()
		close(done)
		wg.Wait()
	}()

	// Setup modeler so we control the contents of all 3 stores: assumed,
	// scheduled and queued
	scheduledPodStore := cache.NewStore(cache.MetaNamespaceKeyFunc)
	scheduledPodLister := &cache.StoreToPodLister{Store: scheduledPodStore}

	queuedPodStore := cache.NewFIFO(cache.MetaNamespaceKeyFunc)
	queuedPodLister := &cache.StoreToPodLister{Store: queuedPodStore}

	modeler := NewSimpleModeler(queuedPodLister, scheduledPodLister)

	// Create a fake clock used to timestamp entries and calculate ttl. Nothing
	// will expire till we flip to something older than the ttl, at which point
	// all entries inserted with fakeTime will expire.
	ttl := 30 * time.Second
	fakeTime := time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC)
	fakeClock := util.NewFakeClock(fakeTime)
	ttlPolicy := &cache.TTLPolicy{Ttl: ttl, Clock: fakeClock}
	assumedPodsStore := cache.NewFakeExpirationStore(
		cache.MetaNamespaceKeyFunc, nil, ttlPolicy, fakeClock)
	modeler.assumedPods = &cache.StoreToPodLister{Store: assumedPodsStore}

	// Port is the easiest way to cause a fit predicate failure
	podPort := 8080
	firstPod := podWithPort("foo", "", podPort)

	// Create the scheduler config
	algo := NewGenericScheduler(
		map[string]algorithm.FitPredicate{"PodFitsHostPorts": predicates.PodFitsHostPorts},
		[]algorithm.PriorityConfig{},
		[]algorithm.SchedulerExtender{},
		modeler.PodLister(),
		rand.New(rand.NewSource(time.Now().UnixNano())))

	var gotBinding *api.Binding
	c := &Config{
		Modeler: modeler,
		NodeLister: algorithm.FakeNodeLister(
			api.NodeList{Items: []api.Node{{ObjectMeta: api.ObjectMeta{Name: "machine1"}}}},
		),
		Algorithm: algo,
		Binder: fakeBinder{func(b *api.Binding) error {
			scheduledPodStore.Add(podWithPort(b.Name, b.Target.Name, podPort))
			gotBinding = b
			return nil
		}},
		NextPod: func() *api.Pod {
			return queuedPodStore.Pop().(*api.Pod)
		},
		Error: func(p *api.Pod, err error) {
			t.Errorf("Unexpected error when scheduling pod %+v: %v", p, err)
		},
		Recorder: eventBroadcaster.NewRecorder(api.EventSource{Component: "scheduler"}),
	}

	// First scheduling pass should schedule the pod
	s := New(c)
	called := make(chan struct{})
	events := eventBroadcaster.StartEventWatcher(func(e *api.Event) {
		if e, a := "Scheduled", e.Reason; e != a {
			t.Errorf("expected %v, got %v", e, a)
		}
		close(called)
	})

	queuedPodStore.Add(firstPod)
	// queuedPodStore: [foo:8080]
	// scheduledPodStore: []
	// assumedPods: []

	s.scheduleOne()
	// queuedPodStore: []
	// scheduledPodStore: [foo:8080]
	// assumedPods: [foo:8080]

	pod, exists, _ := scheduledPodStore.GetByKey("foo")
	if !exists {
		t.Errorf("Expected scheduled pod store to contain pod")
	}
	pod, exists, _ = queuedPodStore.GetByKey("foo")
	if exists {
		t.Errorf("Did not expect a queued pod, found %+v", pod)
	}
	pod, exists, _ = assumedPodsStore.GetByKey("foo")
	if !exists {
		t.Errorf("Assumed pod store should contain stale pod")
	}

	expectBind := &api.Binding{
		ObjectMeta: api.ObjectMeta{Name: "foo"},
		Target:     api.ObjectReference{Kind: "Node", Name: "machine1"},
	}
	if ex, ac := expectBind, gotBinding; !reflect.DeepEqual(ex, ac) {
		t.Errorf("Expected exact match on binding: %s", util.ObjectDiff(ex, ac))
	}

	<-called
	events.Stop()

	scheduledPodStore.Delete(pod)
	_, exists, _ = assumedPodsStore.Get(pod)
	if !exists {
		t.Errorf("Expected pod %#v in assumed pod store", pod)
	}

	secondPod := podWithPort("bar", "", podPort)
	queuedPodStore.Add(secondPod)
	// queuedPodStore: [bar:8080]
	// scheduledPodStore: []
	// assumedPods: [foo:8080]

	// Second scheduling pass will fail to schedule if the store hasn't expired
	// the deleted pod. This would normally happen with a timeout.
	//expirationPolicy.NeverExpire = util.NewStringSet()
	fakeClock.Step(ttl + 1)

	called = make(chan struct{})
	events = eventBroadcaster.StartEventWatcher(func(e *api.Event) {
		if e, a := "Scheduled", e.Reason; e != a {
			t.Errorf("expected %v, got %v", e, a)
		}
		close(called)
	})

	s.scheduleOne()

	expectBind = &api.Binding{
		ObjectMeta: api.ObjectMeta{Name: "bar"},
		Target:     api.ObjectReference{Kind: "Node", Name: "machine1"},
	}
	if ex, ac := expectBind, gotBinding; !reflect.DeepEqual(ex, ac) {
		t.Errorf("Expected exact match on binding: %s", util.ObjectDiff(ex, ac))
	}
	<-called
	events.Stop()
}
<|endoftext|>


<commit_before>package leveldb

import (
	"math/rand"
	"runtime"

	"github.com/kezhuw/leveldb/internal/errors"
	"github.com/kezhuw/leveldb/internal/iterator"
	"github.com/kezhuw/leveldb/internal/keys"
	"github.com/kezhuw/leveldb/internal/manifest"
)

type dbIterator struct {
	db       *DB
	base     *manifest.Version
	ucmp     keys.Comparator
	sequence keys.Sequence

	err       error
	status    iterator.Status
	direction iterator.Direction
	iterator  iterator.Iterator

	// The last key we read from underlying iterator, used as
	// current key for this iterator.
	lastKey   []byte
	lastValue []byte
	parsedKey keys.ParsedInternalKey

	rnd         *rand.Rand
	sampleBytes int

	tag        [8]byte // packed sequence and keys.Seek
	keyScratch [256]byte
}

func (it *dbIterator) First() bool {
	it.err = nil
	it.direction = iterator.Forward
	if !it.iterator.First() {
		it.status = iterator.Invalid
		return false
	}
	it.status = iterator.Valid
	return it.findNextEntry(false)
}

func (it *dbIterator) Last() bool {
	it.err = nil
	it.direction = iterator.Reverse
	if !it.iterator.Last() {
		it.err = it.iterator.Err()
		it.status = iterator.Invalid
		return false
	}
	it.status = iterator.Valid
	return it.findPrevEntry()
}

func (it *dbIterator) Seek(key []byte) bool {
	it.err = nil
	it.lastKey = append(it.lastKey[:0], key...)
	it.lastKey = append(it.lastKey, it.tag[:]...)
	if !it.iterator.Seek(it.lastKey) {
		it.err = it.iterator.Err()
		it.status = iterator.Invalid
		return false
	}
	it.status = iterator.Valid
	it.direction = iterator.Forward
	return it.findNextEntry(false)
}

func (it *dbIterator) Next() bool {
	switch {
	case it.status == iterator.Initial:
		return it.First()
	case it.status != iterator.Valid:
		return false
	case it.direction == iterator.Reverse:
		// Two possible situations exist:
		//
		// * Iterator points at the last entry which has a user key less than
		//   current one.
		// * Iterator is invalid, current key has the smallest user key in iterator.
		//
		// In either case, there may be some large sequence incarnations before current
		// key, but it doesn't matter since we will ignore them all.
		//
		// So we can simply advance to the next or first entry, which has same
		// user key as current, to reverse the direction.
		if !iterator.Next(it.iterator) {
			it.err = it.iterator.Err()
			it.status = iterator.Invalid
			return false
		}
		it.direction = iterator.Forward
		fallthrough
	default:
		return it.findNextEntry(true)
	}
}

func (it *dbIterator) Prev() bool {
	switch {
	case it.status == iterator.Initial:
		return it.Last()
	case it.status != iterator.Valid:
		return false
	case it.direction == iterator.Forward:
		// Iterator is pointing at current key which has valid
		// sequence(<= it.sequence) and Kind keys.Value. There is no
		// possibility that previous valid entry has same user key as
		// current, if that, it should shadow current one.
		it.iterator.Prev()
		it.direction = iterator.Reverse
		fallthrough
	default:
		return it.findPrevEntry()
	}
}

func (it *dbIterator) Valid() bool {
	return it.status == iterator.Valid
}

func (it *dbIterator) Key() []byte {
	return it.lastKey
}

func (it *dbIterator) Value() []byte {
	switch it.direction {
	case iterator.Forward:
		return it.iterator.Value()
	default:
		return it.lastValue
	}
}

func (it *dbIterator) Err() error {
	return it.err
}

func (it *dbIterator) finalize() error {
	if it.status == iterator.Closed {
		return nil
	}
	it.status = iterator.Closed
	it.db = nil
	it.base = nil
	it.iterator = nil
	return it.Err()
}

func (it *dbIterator) Release() error {
	runtime.SetFinalizer(it, nil)
	return it.finalize()
}

func (it *dbIterator) findNextEntry(skip bool) bool {
	ikey := &it.parsedKey
	for {
		if !ikey.Parse(it.iterator.Key()) {
			it.err = errors.ErrCorruptInternalKey
			it.status = iterator.Invalid
			return false
		}
		if ikey.Sequence <= it.sequence {
			switch ikey.Kind {
			case keys.Delete:
				it.lastKey = append(it.lastKey[:0], ikey.UserKey...)
				skip = true
			case keys.Value:
				if skip && it.ucmp.Compare(ikey.UserKey, it.lastKey) <= 0 {
					break
				}
				it.lastKey = append(it.lastKey[:0], ikey.UserKey...)
				return true
			default:
				it.err = errors.ErrCorruptInternalKey
				it.status = iterator.Invalid
				return false
			}
		}
		if !it.iterator.Next() {
			it.err = it.iterator.Err()
			it.status = iterator.Invalid
			return false
		}
	}
}

func (it *dbIterator) randomSampleBytes() int {
	return it.rnd.Intn(2 * it.db.options.IterationBytesPerSampleSeek)
}

func (it *dbIterator) parseKey(ikey *keys.ParsedInternalKey) bool {
	key := it.iterator.Key()
	if !ikey.Parse(key) {
		it.err = errors.ErrCorruptInternalKey
		it.status = iterator.Invalid
		return false
	}
	it.sampleBytes -= len(key) + len(it.iterator.Value())
	for it.sampleBytes < 0 {
		it.sampleBytes += it.randomSampleBytes()
		seekOverlapFile := it.base.SeekOverlap(key, nil)
		if seekOverlapFile.FileMeta != nil {
			it.db.tryCompactFile(seekOverlapFile)
		}
	}
	return true
}

func (it *dbIterator) findPrevEntry() bool {
	if !it.iterator.Valid() {
		it.err = it.iterator.Err()
		it.status = iterator.Invalid
		return false
	}
	ikey := &it.parsedKey
	skip := true
	it.status = iterator.Invalid
	for {
		if !ikey.Parse(it.iterator.Key()) {
			return false
		}
		if ikey.Sequence <= it.sequence {
			if !skip && it.ucmp.Compare(ikey.UserKey, it.lastKey) < 0 {
				return true
			}
			switch {
			case ikey.Kind == keys.Delete:
				skip = true
				it.status = iterator.Invalid
			case ikey.Kind == keys.Value:
				skip = false
				it.status = iterator.Valid
				it.lastKey = append(it.lastKey[:0], ikey.UserKey...)
				it.lastValue = append(it.lastValue[:0], it.iterator.Value()...)
			default:
				it.err = errors.ErrCorruptInternalKey
				it.status = iterator.Invalid
				return false
			}
		}
		if !it.iterator.Prev() {
			// If there is an error, current key may be
			// shadowed by previous entry.
			if err := it.iterator.Err(); err != nil {
				it.err = err
				it.status = iterator.Invalid
				return false
			}
			return skip == false
		}
	}
}

func newDBIterator(db *DB, base *manifest.Version, seq keys.Sequence, it iterator.Iterator) iterator.Iterator {
	rnd := rand.New(rand.NewSource(rand.Int63()))
	dbIt := &dbIterator{
		db:       db,
		base:     base,
		ucmp:     db.options.Comparator.UserKeyComparator,
		iterator: it,
		sequence: seq,
		rnd:      rnd,
	}
	runtime.SetFinalizer(dbIt, (*dbIterator).finalize)
	keys.CombineTag(dbIt.tag[:], seq, keys.Seek)
	dbIt.lastKey = dbIt.keyScratch[:0]
	dbIt.sampleBytes = dbIt.randomSampleBytes()
	return dbIt
}
<commit_msg>Sample seeks base on current manifest.Version always<commit_after>package leveldb

import (
	"math/rand"
	"runtime"

	"github.com/kezhuw/leveldb/internal/errors"
	"github.com/kezhuw/leveldb/internal/iterator"
	"github.com/kezhuw/leveldb/internal/keys"
	"github.com/kezhuw/leveldb/internal/manifest"
)

type dbIterator struct {
	db *DB
	// Keep it away from GC, this way level files iterator seeks in
	// wouldn't got deleted due to umount from manifest.
	base     *manifest.Version
	ucmp     keys.Comparator
	sequence keys.Sequence

	err       error
	status    iterator.Status
	direction iterator.Direction
	iterator  iterator.Iterator

	// The last key we read from underlying iterator, used as
	// current key for this iterator.
	lastKey   []byte
	lastValue []byte
	parsedKey keys.ParsedInternalKey

	rnd         *rand.Rand
	sampleBytes int

	tag        [8]byte // packed sequence and keys.Seek
	keyScratch [256]byte
}

func (it *dbIterator) First() bool {
	it.err = nil
	it.direction = iterator.Forward
	if !it.iterator.First() {
		it.status = iterator.Invalid
		return false
	}
	it.status = iterator.Valid
	return it.findNextEntry(false)
}

func (it *dbIterator) Last() bool {
	it.err = nil
	it.direction = iterator.Reverse
	if !it.iterator.Last() {
		it.err = it.iterator.Err()
		it.status = iterator.Invalid
		return false
	}
	it.status = iterator.Valid
	return it.findPrevEntry()
}

func (it *dbIterator) Seek(key []byte) bool {
	it.err = nil
	it.lastKey = append(it.lastKey[:0], key...)
	it.lastKey = append(it.lastKey, it.tag[:]...)
	if !it.iterator.Seek(it.lastKey) {
		it.err = it.iterator.Err()
		it.status = iterator.Invalid
		return false
	}
	it.status = iterator.Valid
	it.direction = iterator.Forward
	return it.findNextEntry(false)
}

func (it *dbIterator) Next() bool {
	switch {
	case it.status == iterator.Initial:
		return it.First()
	case it.status != iterator.Valid:
		return false
	case it.direction == iterator.Reverse:
		// Two possible situations exist:
		//
		// * Iterator points at the last entry which has a user key less than
		//   current one.
		// * Iterator is invalid, current key has the smallest user key in iterator.
		//
		// In either case, there may be some large sequence incarnations before current
		// key, but it doesn't matter since we will ignore them all.
		//
		// So we can simply advance to the next or first entry, which has same
		// user key as current, to reverse the direction.
		if !iterator.Next(it.iterator) {
			it.err = it.iterator.Err()
			it.status = iterator.Invalid
			return false
		}
		it.direction = iterator.Forward
		fallthrough
	default:
		return it.findNextEntry(true)
	}
}

func (it *dbIterator) Prev() bool {
	switch {
	case it.status == iterator.Initial:
		return it.Last()
	case it.status != iterator.Valid:
		return false
	case it.direction == iterator.Forward:
		// Iterator is pointing at current key which has valid
		// sequence(<= it.sequence) and Kind keys.Value. There is no
		// possibility that previous valid entry has same user key as
		// current, if that, it should shadow current one.
		it.iterator.Prev()
		it.direction = iterator.Reverse
		fallthrough
	default:
		return it.findPrevEntry()
	}
}

func (it *dbIterator) Valid() bool {
	return it.status == iterator.Valid
}

func (it *dbIterator) Key() []byte {
	return it.lastKey
}

func (it *dbIterator) Value() []byte {
	switch it.direction {
	case iterator.Forward:
		return it.iterator.Value()
	default:
		return it.lastValue
	}
}

func (it *dbIterator) Err() error {
	return it.err
}

func (it *dbIterator) finalize() error {
	if it.status == iterator.Closed {
		return nil
	}
	it.status = iterator.Closed
	it.db = nil
	it.base = nil
	it.iterator = nil
	return it.Err()
}

func (it *dbIterator) Release() error {
	runtime.SetFinalizer(it, nil)
	return it.finalize()
}

func (it *dbIterator) findNextEntry(skip bool) bool {
	ikey := &it.parsedKey
	for {
		if !ikey.Parse(it.iterator.Key()) {
			it.err = errors.ErrCorruptInternalKey
			it.status = iterator.Invalid
			return false
		}
		if ikey.Sequence <= it.sequence {
			switch ikey.Kind {
			case keys.Delete:
				it.lastKey = append(it.lastKey[:0], ikey.UserKey...)
				skip = true
			case keys.Value:
				if skip && it.ucmp.Compare(ikey.UserKey, it.lastKey) <= 0 {
					break
				}
				it.lastKey = append(it.lastKey[:0], ikey.UserKey...)
				return true
			default:
				it.err = errors.ErrCorruptInternalKey
				it.status = iterator.Invalid
				return false
			}
		}
		if !it.iterator.Next() {
			it.err = it.iterator.Err()
			it.status = iterator.Invalid
			return false
		}
	}
}

func (it *dbIterator) randomSampleBytes() int {
	return it.rnd.Intn(2 * it.db.options.IterationBytesPerSampleSeek)
}

func (it *dbIterator) parseKey(ikey *keys.ParsedInternalKey) bool {
	key := it.iterator.Key()
	if !ikey.Parse(key) {
		it.err = errors.ErrCorruptInternalKey
		it.status = iterator.Invalid
		return false
	}
	it.sampleBytes -= len(key) + len(it.iterator.Value())
	for it.sampleBytes < 0 {
		it.sampleBytes += it.randomSampleBytes()
		seekOverlapFile := it.db.manifest.Version().SeekOverlap(key, nil)
		if seekOverlapFile.FileMeta != nil {
			it.db.tryCompactFile(seekOverlapFile)
		}
	}
	return true
}

func (it *dbIterator) findPrevEntry() bool {
	if !it.iterator.Valid() {
		it.err = it.iterator.Err()
		it.status = iterator.Invalid
		return false
	}
	ikey := &it.parsedKey
	skip := true
	it.status = iterator.Invalid
	for {
		if !ikey.Parse(it.iterator.Key()) {
			return false
		}
		if ikey.Sequence <= it.sequence {
			if !skip && it.ucmp.Compare(ikey.UserKey, it.lastKey) < 0 {
				return true
			}
			switch {
			case ikey.Kind == keys.Delete:
				skip = true
				it.status = iterator.Invalid
			case ikey.Kind == keys.Value:
				skip = false
				it.status = iterator.Valid
				it.lastKey = append(it.lastKey[:0], ikey.UserKey...)
				it.lastValue = append(it.lastValue[:0], it.iterator.Value()...)
			default:
				it.err = errors.ErrCorruptInternalKey
				it.status = iterator.Invalid
				return false
			}
		}
		if !it.iterator.Prev() {
			// If there is an error, current key may be
			// shadowed by previous entry.
			if err := it.iterator.Err(); err != nil {
				it.err = err
				it.status = iterator.Invalid
				return false
			}
			return skip == false
		}
	}
}

func newDBIterator(db *DB, base *manifest.Version, seq keys.Sequence, it iterator.Iterator) iterator.Iterator {
	rnd := rand.New(rand.NewSource(rand.Int63()))
	dbIt := &dbIterator{
		db:       db,
		base:     base,
		ucmp:     db.options.Comparator.UserKeyComparator,
		iterator: it,
		sequence: seq,
		rnd:      rnd,
	}
	runtime.SetFinalizer(dbIt, (*dbIterator).finalize)
	keys.CombineTag(dbIt.tag[:], seq, keys.Seek)
	dbIt.lastKey = dbIt.keyScratch[:0]
	dbIt.sampleBytes = dbIt.randomSampleBytes()
	return dbIt
}
<|endoftext|>


<commit_before>package next

import (
	"database/sql"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"reflect"
	"regexp"
	"runtime"
	"strings"
)

type Mysql struct {
	Db *sql.DB
	Tx *sql.Tx
}

func NewMysql() *Mysql {
	return &Mysql{}
}

func (mysql *Mysql) Open(dsn string) {
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		panic(err.Error())
	}

	mysql.Db = db
	runtime.SetFinalizer(db, nil)
}

func (mysql *Mysql) Ping() error {
	return mysql.Db.Ping()
}

func (mysql *Mysql) Close() {
	mysql.Db.Close()
}

func (mysql *Mysql) Query(args ...interface{}) ([]interface{}, error) {
	// Parse sql
	q, param, err := mysql.parseQuery(args...)
	if err != nil {
		return nil, err
	}

	var stmt *sql.Stmt
	// Prepare
	if mysql.Tx != nil {
		stmt, err = mysql.Tx.Prepare(q)
	} else {
		stmt, err = mysql.Db.Prepare(q)
	}
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	// -------------
	// Scan
	// -------------
	rows, err := stmt.Query(param...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// Get all col name
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	// Scan value
	out := make([]interface{}, 0)
	for rows.Next() {
		rbs := make([]sql.RawBytes, len(cols))
		vals := make([]interface{}, len(cols))
		for i := range vals {
			vals[i] = &rbs[i]
		}

		if err = rows.Scan(vals...); err != nil {
			return nil, err
		}

		// Make map
		r := make(map[string]interface{}, len(cols))
		for i, val := range rbs {
			r[cols[i]] = string(val)
		}
		out = append(out, r)
	}

	return out, nil
}

func (mysql *Mysql) Row(args ...interface{}) (interface{}, error) {
	// Parse sql
	q, param, err := mysql.parseQuery(args...)
	if err != nil {
		return nil, err
	}

	var stmt *sql.Stmt
	// Prepare
	if mysql.Tx != nil {
		stmt, err = mysql.Tx.Prepare(q)
	} else {
		stmt, err = mysql.Db.Prepare(q)
	}
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	// -------------
	// Scan
	// -------------
	rows, err := stmt.Query(param...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// Get all col name
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	// Scan value
	out := make(map[string]interface{}, len(cols))
	if rows.Next() {
		rbs := make([]sql.RawBytes, len(cols))
		vals := make([]interface{}, len(cols))
		for i := range vals {
			vals[i] = &rbs[i]
		}

		if err = rows.Scan(vals...); err != nil {
			return nil, err
		}

		// Make map
		for i, val := range rbs {
			out[cols[i]] = string(val)
		}
	}

	return out, nil
}

// DML for tabel, insert update delete operation.
// usage:
// mysql.Exec("users", "add", &addData)
// mysql.Exec("users", "up", &upData, &where)
// mysql.Exec("users", "del", &where)
func (mysql *Mysql) Exec(table, dml string, args ...interface{}) (interface{}, error) {
	if len(table) == 0 {
		panic("plase set table for exec method")
	}

	switch strings.ToLower(dml); dml {
	case "add":
		return mysql.insert(table, args...)
	case "up":
		return mysql.update(table, args...)
	case "del":
		return mysql.del(table, args...)
	default:
		panic("mysql exec method <add, up, del> can not support: " + dml)
	}
	return nil, nil
}

func (mysql *Mysql) Begin() {
	tx, err := mysql.Db.Begin()
	if err != nil {
		panic(err.Error())
	}
	mysql.Tx = tx
}

func (mysql *Mysql) Rollback() error {
	if mysql.Tx != nil {
		return mysql.Tx.Rollback()
	}
	return nil
}

func (mysql *Mysql) Commit() error {
	if mysql.Tx != nil {
		err := mysql.Tx.Commit()
		mysql.Tx = nil
		return err
	}
	return nil
}

func (mysql *Mysql) parseQuery(args ...interface{}) (string, []interface{}, error) {
	// SQL
	q, ok := (args[0]).(string)
	if !ok {
		panic("first param should be a sql string")
	}

	// No param
	if len(args) == 1 {
		return q, nil, nil
	}

	// Parse sql with mapping
	v := reflect.ValueOf(args[1])
	if v.Kind() != reflect.Ptr || v.Elem().Kind() != reflect.Map {
		panic("param should be a map's pointer")
	}

	param := make([]interface{}, 0)
	re := regexp.MustCompile(":\\w+")
	q = re.ReplaceAllStringFunc(q, func(src string) string {
		param = append(param, v.Elem().MapIndex(reflect.ValueOf(src[1:])).Interface())
		return "?"
	})

	return q, param, nil
}

func (mysql *Mysql) insert(table string, args ...interface{}) (interface{}, error) {
	v := reflect.ValueOf(args[0])
	if v.Kind() != reflect.Ptr || v.Elem().Kind() != reflect.Map {
		panic("param should be a map's pointer")
	}

	// Parse args
	maps := v.Elem().MapKeys()
	feild := make([]string, len(maps))
	bind := make([]string, len(maps))
	param := make([]interface{}, 0)
	for i, k := range maps {
		f := strings.SplitN(k.String(), ":", 2)
		if len(f) > 1 {
			feild[i] = fmt.Sprintf("`%s`", f[0])
			bind[i] = fmt.Sprintf("%s", v.Elem().MapIndex(k).Interface())
		} else {
			feild[i] = fmt.Sprintf("`%s`", k)
			bind[i] = "?"
			param = append(param, v.Elem().MapIndex(k).Interface())
		}
	}
	q := fmt.Sprintf("INSERT INTO `%s`(%s) VALUES(%s);", table, strings.Join(feild, ","), strings.Join(bind, ","))

	// Sql execute
	var stmt *sql.Stmt
	var err error
	// Prepare
	if mysql.Tx != nil {
		stmt, err = mysql.Tx.Prepare(q)
	} else {
		stmt, err = mysql.Db.Prepare(q)
	}
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	res, err := stmt.Exec(param...)
	if err != nil {
		return nil, err
	}

	return res.LastInsertId()
}

func (mysql *Mysql) update(table string, args ...interface{}) (interface{}, error) {
	if len(args) < 2 {
		panic("mysql update method need the data of update and where")
	}

	u := reflect.ValueOf(args[0])
	if u.Kind() != reflect.Ptr || u.Elem().Kind() != reflect.Map {
		panic("updata args should be a map's pointer")
	}
	w := reflect.ValueOf(args[1])
	if w.Kind() != reflect.Ptr || w.Elem().Kind() != reflect.Map {
		panic("where args should be a map's pointer")
	}

	param := make([]interface{}, 0)
	// Parse upate args
	mu := u.Elem().MapKeys()
	up := make([]string, len(mu))
	for i, k := range mu {
		f := strings.SplitN(k.String(), ":", 2)
		if len(f) > 1 {
			up[i] = fmt.Sprintf("`%s`%s%s", f[0], f[1], u.Elem().MapIndex(k).Interface())
		} else {
			up[i] = fmt.Sprintf("`%s`=?", k)
			param = append(param, u.Elem().MapIndex(k).Interface())
		}
	}
	// Parse upate args
	mw := w.Elem().MapKeys()
	where := make([]string, len(mw))
	for i, k := range mw {
		f := strings.SplitN(k.String(), ":", 2)
		if len(f) > 1 {
			where[i] = fmt.Sprintf("`%s`%s%s", f[0], f[1], w.Elem().MapIndex(k).Interface())
		} else {
			where[i] = fmt.Sprintf("`%s`=?", k)
			param = append(param, w.Elem().MapIndex(k).Interface())
		}
	}
	q := fmt.Sprintf("UPDATE `%s` SET %s WHERE %s;", table, strings.Join(up, ","), strings.Join(where, " AND "))

	// Sql execute
	var stmt *sql.Stmt
	var err error
	// Prepare
	if mysql.Tx != nil {
		stmt, err = mysql.Tx.Prepare(q)
	} else {
		stmt, err = mysql.Db.Prepare(q)
	}
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	res, err := stmt.Exec(param...)
	if err != nil {
		return nil, err
	}

	return res.RowsAffected()
}

func (mysql *Mysql) del(table string, args ...interface{}) (interface{}, error) {
	w := reflect.ValueOf(args[0])
	if w.Kind() != reflect.Ptr || w.Elem().Kind() != reflect.Map {
		panic("where args should be a map's pointer")
	}

	mw := w.Elem().MapKeys()
	param := make([]interface{}, 0)
	// Parse where args
	where := make([]string, len(mw))
	for i, k := range mw {
		f := strings.SplitN(k.String(), ":", 2)
		if len(f) > 1 {
			where[i] = fmt.Sprintf("`%s`%s%s", f[0], f[1], w.Elem().MapIndex(k).Interface())
		} else {
			where[i] = fmt.Sprintf("`%s`=?", k)
			param = append(param, w.Elem().MapIndex(k).Interface())
		}
	}
	q := fmt.Sprintf("DELETE FROM `%s` WHERE %s;", table, strings.Join(where, " AND "))

	// Sql execute
	var stmt *sql.Stmt
	var err error
	// Prepare
	if mysql.Tx != nil {
		stmt, err = mysql.Tx.Prepare(q)
	} else {
		stmt, err = mysql.Db.Prepare(q)
	}
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	res, err := stmt.Exec(param...)
	if err != nil {
		return nil, err
	}

	return res.RowsAffected()
}
<commit_msg>fix commit issue<commit_after>package next

import (
	"database/sql"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"reflect"
	"regexp"
	"runtime"
	"strings"
)

type Mysql struct {
	Db *sql.DB
	Tx *sql.Tx
}

func NewMysql() *Mysql {
	return &Mysql{}
}

func (mysql *Mysql) Open(dsn string) {
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		panic(err.Error())
	}

	mysql.Db = db
	runtime.SetFinalizer(db, nil)
}

func (mysql *Mysql) Ping() error {
	return mysql.Db.Ping()
}

func (mysql *Mysql) Close() {
	mysql.Db.Close()
}

func (mysql *Mysql) Query(args ...interface{}) ([]interface{}, error) {
	// Parse sql
	q, param, err := mysql.parseQuery(args...)
	if err != nil {
		return nil, err
	}

	var stmt *sql.Stmt
	// Prepare
	if mysql.Tx != nil {
		stmt, err = mysql.Tx.Prepare(q)
	} else {
		stmt, err = mysql.Db.Prepare(q)
	}
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	// -------------
	// Scan
	// -------------
	rows, err := stmt.Query(param...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// Get all col name
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	// Scan value
	out := make([]interface{}, 0)
	for rows.Next() {
		rbs := make([]sql.RawBytes, len(cols))
		vals := make([]interface{}, len(cols))
		for i := range vals {
			vals[i] = &rbs[i]
		}

		if err = rows.Scan(vals...); err != nil {
			return nil, err
		}

		// Make map
		r := make(map[string]interface{}, len(cols))
		for i, val := range rbs {
			r[cols[i]] = string(val)
		}
		out = append(out, r)
	}

	return out, nil
}

func (mysql *Mysql) Row(args ...interface{}) (interface{}, error) {
	// Parse sql
	q, param, err := mysql.parseQuery(args...)
	if err != nil {
		return nil, err
	}

	var stmt *sql.Stmt
	// Prepare
	if mysql.Tx != nil {
		stmt, err = mysql.Tx.Prepare(q)
	} else {
		stmt, err = mysql.Db.Prepare(q)
	}
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	// -------------
	// Scan
	// -------------
	rows, err := stmt.Query(param...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// Get all col name
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}

	// Scan value
	out := make(map[string]interface{}, len(cols))
	if rows.Next() {
		rbs := make([]sql.RawBytes, len(cols))
		vals := make([]interface{}, len(cols))
		for i := range vals {
			vals[i] = &rbs[i]
		}

		if err = rows.Scan(vals...); err != nil {
			return nil, err
		}

		// Make map
		for i, val := range rbs {
			out[cols[i]] = string(val)
		}
	}

	return out, nil
}

// DML for tabel, insert update delete operation.
// usage:
// mysql.Exec("users", "add", &addData)
// mysql.Exec("users", "up", &upData, &where)
// mysql.Exec("users", "del", &where)
func (mysql *Mysql) Exec(table, dml string, args ...interface{}) (interface{}, error) {
	if len(table) == 0 {
		panic("plase set table for exec method")
	}

	switch strings.ToLower(dml); dml {
	case "add":
		return mysql.insert(table, args...)
	case "up":
		return mysql.update(table, args...)
	case "del":
		return mysql.del(table, args...)
	default:
		panic("mysql exec method <add, up, del> can not support: " + dml)
	}
	return nil, nil
}

func (mysql *Mysql) Begin() {
	tx, err := mysql.Db.Begin()
	if err != nil {
		panic(err.Error())
	}
	mysql.Tx = tx
}

func (mysql *Mysql) Rollback() error {
	if mysql.Tx != nil {
		err := mysql.Tx.Rollback()
		mysql.Tx = nil
		return err
	}
	return nil
}

func (mysql *Mysql) Commit() error {
	if mysql.Tx != nil {
		err := mysql.Tx.Commit()
		mysql.Tx = nil
		return err
	}
	return nil
}

func (mysql *Mysql) parseQuery(args ...interface{}) (string, []interface{}, error) {
	// SQL
	q, ok := (args[0]).(string)
	if !ok {
		panic("first param should be a sql string")
	}

	// No param
	if len(args) == 1 {
		return q, nil, nil
	}

	// Parse sql with mapping
	v := reflect.ValueOf(args[1])
	if v.Kind() != reflect.Ptr || v.Elem().Kind() != reflect.Map {
		panic("param should be a map's pointer")
	}

	param := make([]interface{}, 0)
	re := regexp.MustCompile(":\\w+")
	q = re.ReplaceAllStringFunc(q, func(src string) string {
		param = append(param, v.Elem().MapIndex(reflect.ValueOf(src[1:])).Interface())
		return "?"
	})

	return q, param, nil
}

func (mysql *Mysql) insert(table string, args ...interface{}) (interface{}, error) {
	v := reflect.ValueOf(args[0])
	if v.Kind() != reflect.Ptr || v.Elem().Kind() != reflect.Map {
		panic("param should be a map's pointer")
	}

	// Parse args
	maps := v.Elem().MapKeys()
	feild := make([]string, len(maps))
	bind := make([]string, len(maps))
	param := make([]interface{}, 0)
	for i, k := range maps {
		f := strings.SplitN(k.String(), ":", 2)
		if len(f) > 1 {
			feild[i] = fmt.Sprintf("`%s`", f[0])
			bind[i] = fmt.Sprintf("%s", v.Elem().MapIndex(k).Interface())
		} else {
			feild[i] = fmt.Sprintf("`%s`", k)
			bind[i] = "?"
			param = append(param, v.Elem().MapIndex(k).Interface())
		}
	}
	q := fmt.Sprintf("INSERT INTO `%s`(%s) VALUES(%s);", table, strings.Join(feild, ","), strings.Join(bind, ","))

	// Sql execute
	var stmt *sql.Stmt
	var err error
	// Prepare
	if mysql.Tx != nil {
		stmt, err = mysql.Tx.Prepare(q)
	} else {
		stmt, err = mysql.Db.Prepare(q)
	}
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	res, err := stmt.Exec(param...)
	if err != nil {
		return nil, err
	}

	return res.LastInsertId()
}

func (mysql *Mysql) update(table string, args ...interface{}) (interface{}, error) {
	if len(args) < 2 {
		panic("mysql update method need the data of update and where")
	}

	u := reflect.ValueOf(args[0])
	if u.Kind() != reflect.Ptr || u.Elem().Kind() != reflect.Map {
		panic("updata args should be a map's pointer")
	}
	w := reflect.ValueOf(args[1])
	if w.Kind() != reflect.Ptr || w.Elem().Kind() != reflect.Map {
		panic("where args should be a map's pointer")
	}

	param := make([]interface{}, 0)
	// Parse upate args
	mu := u.Elem().MapKeys()
	up := make([]string, len(mu))
	for i, k := range mu {
		f := strings.SplitN(k.String(), ":", 2)
		if len(f) > 1 {
			up[i] = fmt.Sprintf("`%s`%s%s", f[0], f[1], u.Elem().MapIndex(k).Interface())
		} else {
			up[i] = fmt.Sprintf("`%s`=?", k)
			param = append(param, u.Elem().MapIndex(k).Interface())
		}
	}
	// Parse upate args
	mw := w.Elem().MapKeys()
	where := make([]string, len(mw))
	for i, k := range mw {
		f := strings.SplitN(k.String(), ":", 2)
		if len(f) > 1 {
			where[i] = fmt.Sprintf("`%s`%s%s", f[0], f[1], w.Elem().MapIndex(k).Interface())
		} else {
			where[i] = fmt.Sprintf("`%s`=?", k)
			param = append(param, w.Elem().MapIndex(k).Interface())
		}
	}
	q := fmt.Sprintf("UPDATE `%s` SET %s WHERE %s;", table, strings.Join(up, ","), strings.Join(where, " AND "))

	// Sql execute
	var stmt *sql.Stmt
	var err error
	// Prepare
	if mysql.Tx != nil {
		stmt, err = mysql.Tx.Prepare(q)
	} else {
		stmt, err = mysql.Db.Prepare(q)
	}
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	res, err := stmt.Exec(param...)
	if err != nil {
		return nil, err
	}

	return res.RowsAffected()
}

func (mysql *Mysql) del(table string, args ...interface{}) (interface{}, error) {
	w := reflect.ValueOf(args[0])
	if w.Kind() != reflect.Ptr || w.Elem().Kind() != reflect.Map {
		panic("where args should be a map's pointer")
	}

	mw := w.Elem().MapKeys()
	param := make([]interface{}, 0)
	// Parse where args
	where := make([]string, len(mw))
	for i, k := range mw {
		f := strings.SplitN(k.String(), ":", 2)
		if len(f) > 1 {
			where[i] = fmt.Sprintf("`%s`%s%s", f[0], f[1], w.Elem().MapIndex(k).Interface())
		} else {
			where[i] = fmt.Sprintf("`%s`=?", k)
			param = append(param, w.Elem().MapIndex(k).Interface())
		}
	}
	q := fmt.Sprintf("DELETE FROM `%s` WHERE %s;", table, strings.Join(where, " AND "))

	// Sql execute
	var stmt *sql.Stmt
	var err error
	// Prepare
	if mysql.Tx != nil {
		stmt, err = mysql.Tx.Prepare(q)
	} else {
		stmt, err = mysql.Db.Prepare(q)
	}
	if err != nil {
		return nil, err
	}
	defer stmt.Close()

	res, err := stmt.Exec(param...)
	if err != nil {
		return nil, err
	}

	return res.RowsAffected()
}
<|endoftext|>


<commit_before>//Package filesystem provides a filesystem cache implementation for onecache
package filesystem

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"github.com/adelowo/onecache"
)

const (
	defaultFilePerm          os.FileMode = 0666
	defaultDirectoryFilePerm             = 0755
)

func createDirectory(dir string) error {

	return os.MkdirAll(dir, defaultDirectoryFilePerm)
}

type FSStore struct {
	baseDir string
}

//Returns an initialized Filesystem Cache
//If a non-existent directory is passed, it would be created automatically.
//This function Panics if the directory could not be created
func MustNewFSStore(baseDir string) *FSStore {

	_, err := os.Stat(baseDir)

	if err != nil { //Directory does not exist..Let's create it
		if err := createDirectory(baseDir); err != nil {
			panic(fmt.Errorf("Base directory could not be created : %s", err))
		}
	}

	return &FSStore{baseDir}
}

func (fs *FSStore) Set(key string, data interface{}, expiresAt time.Duration) error {

	path := fs.getFilePathFor(key)

	if err := os.MkdirAll(filepath.Dir(path), defaultDirectoryFilePerm); err != nil {
		return err
	}

	i := &onecache.Item{ExpiresAt: time.Now().Add(expiresAt), Data: data}

	b, err := i.ToBytes()

	if err != nil {
		return err
	}

	if err = ioutil.WriteFile(path, b, defaultFilePerm); err != nil {
		return err
	}

	return nil
}

//Fetches a cache key.
//This runs garbage collection on the key if necessary
func (fs *FSStore) Get(key string) (interface{}, error) {

	b, err := ioutil.ReadFile(fs.getFilePathFor(key))

	if err != nil {
		return nil, err
	}

	i, err := onecache.BytesToItem(b)

	if i.IsExpired() {
		fs.Delete(key)
		return nil, onecache.ErrCacheMiss
	}

	return i.Data, nil
}

//Removes a file (cached item) from the disk
func (fs *FSStore) Delete(key string) error {
	return os.RemoveAll(fs.getFilePathFor(key))
}

//Gets a unique path for a cache key.
//This is going to be a directory 3 level deep. Something like "basedir/33/rr/33/hash"
func (fs *FSStore) getFilePathFor(key string) string {
	hashSum := md5.Sum([]byte(key))

	hashSumAsString := hex.EncodeToString(hashSum[:])

	return filepath.Join(fs.baseDir,
		string(hashSumAsString[0:2]),
		string(hashSumAsString[2:4]),
		string(hashSumAsString[4:6]), hashSumAsString)
}
<commit_msg>Implemented the Flush Method<commit_after>//Package filesystem provides a filesystem cache implementation for onecache
package filesystem

import (
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"time"

	"github.com/adelowo/onecache"
)

const (
	defaultFilePerm          os.FileMode = 0666
	defaultDirectoryFilePerm             = 0755
)

func createDirectory(dir string) error {

	return os.MkdirAll(dir, defaultDirectoryFilePerm)
}

type FSStore struct {
	baseDir string
}

//Returns an initialized Filesystem Cache
//If a non-existent directory is passed, it would be created automatically.
//This function Panics if the directory could not be created
func MustNewFSStore(baseDir string) *FSStore {

	_, err := os.Stat(baseDir)

	if err != nil { //Directory does not exist..Let's create it
		if err := createDirectory(baseDir); err != nil {
			panic(fmt.Errorf("Base directory could not be created : %s", err))
		}
	}

	return &FSStore{baseDir}
}

func (fs *FSStore) Set(key string, data interface{}, expiresAt time.Duration) error {

	path := fs.getFilePathFor(key)

	if err := os.MkdirAll(filepath.Dir(path), defaultDirectoryFilePerm); err != nil {
		return err
	}

	i := &onecache.Item{ExpiresAt: time.Now().Add(expiresAt), Data: data}

	b, err := i.ToBytes()

	if err != nil {
		return err
	}

	if err = ioutil.WriteFile(path, b, defaultFilePerm); err != nil {
		return err
	}

	return nil
}

//Fetches a cache key.
//This runs garbage collection on the key if necessary
func (fs *FSStore) Get(key string) (interface{}, error) {

	b, err := ioutil.ReadFile(fs.getFilePathFor(key))

	if err != nil {
		return nil, err
	}

	i, err := onecache.BytesToItem(b)

	if i.IsExpired() {
		fs.Delete(key)
		return nil, onecache.ErrCacheMiss
	}

	return i.Data, nil
}

//Removes a file (cached item) from the disk
func (fs *FSStore) Delete(key string) error {
	return os.RemoveAll(fs.getFilePathFor(key))
}

//Cleans up the entire cache
func (fs *FSStore) Flush() error {
	return os.RemoveAll(fs.baseDir)
}

//Gets a unique path for a cache key.
//This is going to be a directory 3 level deep. Something like "basedir/33/rr/33/hash"
func (fs *FSStore) getFilePathFor(key string) string {
	hashSum := md5.Sum([]byte(key))

	hashSumAsString := hex.EncodeToString(hashSum[:])

	return filepath.Join(fs.baseDir,
		string(hashSumAsString[0:2]),
		string(hashSumAsString[2:4]),
		string(hashSumAsString[4:6]), hashSumAsString)
}
<|endoftext|>


<commit_before>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * rpc/method.go                                          *
 *                                                        *
 * hprose method manager for Go.                          *
 *                                                        *
 * LastModified: Oct 17, 2016                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package rpc

import (
	"reflect"
	"strings"
	"sync"
)

// Options is the options of the published service method
type Options struct {
	Mode           ResultMode
	Simple         bool
	Oneway         bool
	NameSpace      string
	JSONCompatible bool
}

// Method is the published service method
type Method struct {
	Function reflect.Value
	Options
}

// methodManager manages published service methods
type methodManager struct {
	MethodNames   []string
	RemoteMethods map[string]*Method
	sync.RWMutex
}

func (mm *methodManager) initMethodManager() {
	mm.MethodNames = make([]string, 0, 64)
	mm.RemoteMethods = make(map[string]*Method)
}

// AddFunction publish a func or bound method
// name is the method name
// function is a func or bound method
// options includes Mode, Simple, Oneway and NameSpace
func (mm *methodManager) AddFunction(
	name string, function interface{}, options Options) {
	if name == "" {
		panic("name can't be empty")
	}
	if function == nil {
		panic("function can't be nil")
	}
	f, ok := function.(reflect.Value)
	if !ok {
		f = reflect.ValueOf(function)
	}
	if f.Kind() != reflect.Func {
		panic("function must be func or bound method")
	}
	if options.NameSpace != "" && name != "*" {
		name = options.NameSpace + "_" + name
	}
	mm.Lock()
	mm.MethodNames = append(mm.MethodNames, name)
	mm.RemoteMethods[strings.ToLower(name)] = &Method{f, options}
	mm.Unlock()
}

// AddFunctions is used for batch publishing service method
func (mm *methodManager) AddFunctions(
	names []string, functions []interface{}, options Options) {
	count := len(names)
	if count != len(functions) {
		panic("names and functions must have the same length")
	}
	for i := 0; i < count; i++ {
		mm.AddFunction(names[i], functions[i], options)
	}
}

// AddMethod is used for publishing a method on the obj with an alias
func (mm *methodManager) AddMethod(
	name string, obj interface{}, options Options, alias ...string) {
	if obj == nil {
		panic("obj can't be nil")
	}
	f := reflect.ValueOf(obj).MethodByName(name)
	if len(alias) == 1 && alias[0] != "" {
		name = alias[0]
	}
	if f.CanInterface() {
		mm.AddFunction(name, f, options)
	}
}

// AddMethods is used for batch publishing methods on the obj with aliases
func (mm *methodManager) AddMethods(
	names []string, obj interface{}, options Options, aliases ...[]string) {
	if obj == nil {
		panic("obj can't be nil")
	}
	count := len(names)
	if len(aliases) == 1 {
		if len(aliases[0]) != count {
			panic("names and aliases must have the same length")
		}
		for i := 0; i < count; i++ {
			mm.AddMethod(names[i], obj, options, aliases[0][i])
		}
		return
	}
	for i := 0; i < count; i++ {
		mm.AddMethod(names[i], obj, options)
	}
}

func (mm *methodManager) addMethods(
	v reflect.Value, t reflect.Type, options Options) {
	n := t.NumMethod()
	for i := 0; i < n; i++ {
		name := t.Method(i).Name
		method := v.Method(i)
		if method.CanInterface() {
			mm.AddFunction(name, method, options)
		}
	}
}

func getPtrTo(v reflect.Value, t reflect.Type) (reflect.Value, reflect.Type) {
	for ; t.Kind() == reflect.Ptr && !v.IsNil(); v = v.Elem() {
		t = t.Elem()
	}
	return v, t
}

func (mm *methodManager) addFuncField(
	v reflect.Value, t reflect.Type, i int, options Options) {
	f := v.Field(i)
	name := t.Field(i).Name
	if !f.CanInterface() || !f.IsValid() {
		return
	}
	f, _ = getPtrTo(f, f.Type())
	if f.Kind() == reflect.Func && !f.IsNil() {
		mm.AddFunction(name, f, options)
	}
}

func (mm *methodManager) recursiveAddFuncFields(
	v reflect.Value, t reflect.Type, i int, options Options) {
	f := v.Field(i)
	fs := t.Field(i)
	name := fs.Name
	if !f.CanInterface() || !f.IsValid() {
		return
	}
	f, _ = getPtrTo(f, f.Type())
	if f.Kind() == reflect.Func && !f.IsNil() {
		mm.AddFunction(name, f, options)
		return
	}
	if f.Kind() != reflect.Struct {
		return
	}
	if fs.Anonymous {
		mm.AddAllMethods(f.Interface(), options)
	} else {
		newOptions := options
		if newOptions.NameSpace == "" {
			newOptions.NameSpace = name
		} else {
			newOptions.NameSpace += "_" + name
		}
		mm.AddAllMethods(f.Interface(), newOptions)
	}
}

type addFuncFunc func(
	v reflect.Value,
	t reflect.Type,
	i int,
	options Options)

func (mm *methodManager) addInstanceMethods(
	obj interface{}, options Options, addFunc addFuncFunc) {
	if obj == nil {
		panic("obj can't be nil")
	}
	v := reflect.ValueOf(obj)
	t := v.Type()
	mm.addMethods(v, t, options)
	v, t = getPtrTo(v, t)
	if t.Kind() == reflect.Struct {
		n := t.NumField()
		for i := 0; i < n; i++ {
			addFunc(v, t, i, options)
		}
	}
}

// AddInstanceMethods is used for publishing all the public methods and func fields with options.
func (mm *methodManager) AddInstanceMethods(
	obj interface{}, options Options) {
	mm.addInstanceMethods(obj, options, mm.addFuncField)
}

// AddAllMethods will publish all methods and non-nil function fields on the
// obj self and on its anonymous or non-anonymous struct fields (or pointer to
// pointer ... to pointer struct fields). This is a recursive operation.
// So it's a pit, if you do not know what you are doing, do not step on.
func (mm *methodManager) AddAllMethods(
	obj interface{}, options Options) {
	mm.addInstanceMethods(obj, options, mm.recursiveAddFuncFields)
}

// MissingMethod is missing method
type MissingMethod func(name string, args []reflect.Value, context Context) (result []reflect.Value)

// AddMissingMethod is used for publishing a method,
// all methods not explicitly published will be redirected to this method.
func (mm *methodManager) AddMissingMethod(
	method MissingMethod, options Options) {
	mm.AddFunction("*", method, options)
}

// AddNetRPCMethods is used for publishing methods defined for net/rpc.
//
// For example:
//
//		package server
//
//		type Args struct {
//			A, B int
//		}
//
//		type Quotient struct {
//			Quo, Rem int
//		}
//
//		type Arith int
//
//		func (t *Arith) Multiply(args *Args, reply *int) error {
//			*reply = args.A * args.B
//			return nil
//		}
//
//		func (t *Arith) Divide(args *Args, quo *Quotient) error {
//			if args.B == 0 {
//				return errors.New("divide by zero")
//			}
//			quo.Quo = args.A / args.B
//			quo.Rem = args.A % args.B
//			return nil
//		}
//
// You can publish it in hprose like this:
//
//		service := rpc.NewHTTPService()
//		service.AddNetRPCMethods(new(Arith), rpc.Options{})
//		http.ListenAndServe(":8080", service)
//
//
// Then define stub struct in client:
//
//		type Stub struct {
//			// Synchronous call
//			Multiply func(args *Args) int
//			// Asynchronous call
//			Divide func(func(*Quotient, error), *Args)
//		}
//
// Then it can make a remote call:
//
//		client := rpc.NewClient("http://127.0.0.1:8080")
//		var stub *Stub
//		client.UseService(&stub)
//		fmt.Println(stub.Multiply(&Args{8, 7}))
//		stub.Divide(func(result *Quotient, err error) {
//			if err != nil {
//				log.Fatal("arith error:", err)
//			} else {
//				fmt.Println(result.Quo, result.Rem)
//			}
//		}, &Args{8, 7})
//
// You can also call it in other languages.
//
// Here is a simple example in JavaScript:
//
//		// The method name is case-insensitive
//		var methods = ['multiply', 'divide'];
//		var client = hprose.Client.create('http://127.0.0.1:8080/', methods);
//		// The first letter of the field name is convert to lowercase.
//      client.multiply({a:8, b:7}).then(function(product) {
//			return client.divide({a:product, b:6});
//		}).then(function(quotient) {
//			// The first letter of the field name is convert to lowercase.
//			console.log(quotient.quo, quotient.rem);
//		}).catch(function(err) {
//			// The returned err is here
//			console.error(err);
//		});
//
func (mm *methodManager) AddNetRPCMethods(rcvr interface{}, options Options) {
	if rcvr == nil {
		panic("rcvr can't be nil")
	}
	v := reflect.ValueOf(rcvr)
	t := v.Type()
	n := t.NumMethod()
	for i := 0; i < n; i++ {
		name := t.Method(i).Name
		method := v.Method(i)
		if method.CanInterface() {
			mm.addNetRPCMethod(name, method, options)
		}
	}
}

func (mm *methodManager) addNetRPCMethod(
	name string, method reflect.Value, options Options) {
	ft := method.Type()
	if ft.NumIn() != 2 || ft.IsVariadic() {
		panic("the method " + name + " must has two arguments")
	}
	if ft.In(1).Kind() != reflect.Ptr {
		panic("the second argument of method " + name + " must be a pointer")
	}
	if ft.NumOut() != 1 || ft.Out(0) != errorType {
		panic("the result type of method " + name + " must be error")
	}
	argsType := ft.In(0)
	resultType := ft.In(1)
	in := []reflect.Type{argsType}
	out := []reflect.Type{resultType, errorType}
	newft := reflect.FuncOf(in, out, false)
	newMethod := reflect.MakeFunc(newft, func(
		args []reflect.Value) (results []reflect.Value) {
		result := reflect.New(resultType.Elem())
		in := []reflect.Value{args[0], result}
		err := method.Call(in)[0]
		results = []reflect.Value{result, err}
		return
	})
	mm.AddFunction(name, newMethod, options)
}

// Remove the published func or method by name
func (mm *methodManager) Remove(name string) {
	name = strings.ToLower(name)
	mm.Lock()
	n := len(mm.MethodNames)
	for i := 0; i < n; i++ {
		if strings.ToLower(mm.MethodNames[i]) == name {
			mm.MethodNames = append(mm.MethodNames[:i], mm.MethodNames[i+1:]...)
			break
		}
	}
	delete(mm.RemoteMethods, name)
	mm.Unlock()
}
<commit_msg>Improved addNetRPCMethod<commit_after>/**********************************************************\
|                                                          |
|                          hprose                          |
|                                                          |
| Official WebSite: http://www.hprose.com/                 |
|                   http://www.hprose.org/                 |
|                                                          |
\**********************************************************/
/**********************************************************\
 *                                                        *
 * rpc/method.go                                          *
 *                                                        *
 * hprose method manager for Go.                          *
 *                                                        *
 * LastModified: Oct 17, 2016                             *
 * Author: Ma Bingyao <andot@hprose.com>                  *
 *                                                        *
\**********************************************************/

package rpc

import (
	"reflect"
	"strings"
	"sync"
)

// Options is the options of the published service method
type Options struct {
	Mode           ResultMode
	Simple         bool
	Oneway         bool
	NameSpace      string
	JSONCompatible bool
}

// Method is the published service method
type Method struct {
	Function reflect.Value
	Options
}

// methodManager manages published service methods
type methodManager struct {
	MethodNames   []string
	RemoteMethods map[string]*Method
	sync.RWMutex
}

func (mm *methodManager) initMethodManager() {
	mm.MethodNames = make([]string, 0, 64)
	mm.RemoteMethods = make(map[string]*Method)
}

// AddFunction publish a func or bound method
// name is the method name
// function is a func or bound method
// options includes Mode, Simple, Oneway and NameSpace
func (mm *methodManager) AddFunction(
	name string, function interface{}, options Options) {
	if name == "" {
		panic("name can't be empty")
	}
	if function == nil {
		panic("function can't be nil")
	}
	f, ok := function.(reflect.Value)
	if !ok {
		f = reflect.ValueOf(function)
	}
	if f.Kind() != reflect.Func {
		panic("function must be func or bound method")
	}
	if options.NameSpace != "" && name != "*" {
		name = options.NameSpace + "_" + name
	}
	mm.Lock()
	mm.MethodNames = append(mm.MethodNames, name)
	mm.RemoteMethods[strings.ToLower(name)] = &Method{f, options}
	mm.Unlock()
}

// AddFunctions is used for batch publishing service method
func (mm *methodManager) AddFunctions(
	names []string, functions []interface{}, options Options) {
	count := len(names)
	if count != len(functions) {
		panic("names and functions must have the same length")
	}
	for i := 0; i < count; i++ {
		mm.AddFunction(names[i], functions[i], options)
	}
}

// AddMethod is used for publishing a method on the obj with an alias
func (mm *methodManager) AddMethod(
	name string, obj interface{}, options Options, alias ...string) {
	if obj == nil {
		panic("obj can't be nil")
	}
	f := reflect.ValueOf(obj).MethodByName(name)
	if len(alias) == 1 && alias[0] != "" {
		name = alias[0]
	}
	if f.CanInterface() {
		mm.AddFunction(name, f, options)
	}
}

// AddMethods is used for batch publishing methods on the obj with aliases
func (mm *methodManager) AddMethods(
	names []string, obj interface{}, options Options, aliases ...[]string) {
	if obj == nil {
		panic("obj can't be nil")
	}
	count := len(names)
	if len(aliases) == 1 {
		if len(aliases[0]) != count {
			panic("names and aliases must have the same length")
		}
		for i := 0; i < count; i++ {
			mm.AddMethod(names[i], obj, options, aliases[0][i])
		}
		return
	}
	for i := 0; i < count; i++ {
		mm.AddMethod(names[i], obj, options)
	}
}

func (mm *methodManager) addMethods(
	v reflect.Value, t reflect.Type, options Options) {
	n := t.NumMethod()
	for i := 0; i < n; i++ {
		name := t.Method(i).Name
		method := v.Method(i)
		if method.CanInterface() {
			mm.AddFunction(name, method, options)
		}
	}
}

func getPtrTo(v reflect.Value, t reflect.Type) (reflect.Value, reflect.Type) {
	for ; t.Kind() == reflect.Ptr && !v.IsNil(); v = v.Elem() {
		t = t.Elem()
	}
	return v, t
}

func (mm *methodManager) addFuncField(
	v reflect.Value, t reflect.Type, i int, options Options) {
	f := v.Field(i)
	name := t.Field(i).Name
	if !f.CanInterface() || !f.IsValid() {
		return
	}
	f, _ = getPtrTo(f, f.Type())
	if f.Kind() == reflect.Func && !f.IsNil() {
		mm.AddFunction(name, f, options)
	}
}

func (mm *methodManager) recursiveAddFuncFields(
	v reflect.Value, t reflect.Type, i int, options Options) {
	f := v.Field(i)
	fs := t.Field(i)
	name := fs.Name
	if !f.CanInterface() || !f.IsValid() {
		return
	}
	f, _ = getPtrTo(f, f.Type())
	if f.Kind() == reflect.Func && !f.IsNil() {
		mm.AddFunction(name, f, options)
		return
	}
	if f.Kind() != reflect.Struct {
		return
	}
	if fs.Anonymous {
		mm.AddAllMethods(f.Interface(), options)
	} else {
		newOptions := options
		if newOptions.NameSpace == "" {
			newOptions.NameSpace = name
		} else {
			newOptions.NameSpace += "_" + name
		}
		mm.AddAllMethods(f.Interface(), newOptions)
	}
}

type addFuncFunc func(
	v reflect.Value,
	t reflect.Type,
	i int,
	options Options)

func (mm *methodManager) addInstanceMethods(
	obj interface{}, options Options, addFunc addFuncFunc) {
	if obj == nil {
		panic("obj can't be nil")
	}
	v := reflect.ValueOf(obj)
	t := v.Type()
	mm.addMethods(v, t, options)
	v, t = getPtrTo(v, t)
	if t.Kind() == reflect.Struct {
		n := t.NumField()
		for i := 0; i < n; i++ {
			addFunc(v, t, i, options)
		}
	}
}

// AddInstanceMethods is used for publishing all the public methods and func fields with options.
func (mm *methodManager) AddInstanceMethods(
	obj interface{}, options Options) {
	mm.addInstanceMethods(obj, options, mm.addFuncField)
}

// AddAllMethods will publish all methods and non-nil function fields on the
// obj self and on its anonymous or non-anonymous struct fields (or pointer to
// pointer ... to pointer struct fields). This is a recursive operation.
// So it's a pit, if you do not know what you are doing, do not step on.
func (mm *methodManager) AddAllMethods(
	obj interface{}, options Options) {
	mm.addInstanceMethods(obj, options, mm.recursiveAddFuncFields)
}

// MissingMethod is missing method
type MissingMethod func(name string, args []reflect.Value, context Context) (result []reflect.Value)

// AddMissingMethod is used for publishing a method,
// all methods not explicitly published will be redirected to this method.
func (mm *methodManager) AddMissingMethod(
	method MissingMethod, options Options) {
	mm.AddFunction("*", method, options)
}

// AddNetRPCMethods is used for publishing methods defined for net/rpc.
//
// For example:
//
//		package server
//
//		type Args struct {
//			A, B int
//		}
//
//		type Quotient struct {
//			Quo, Rem int
//		}
//
//		type Arith int
//
//		func (t *Arith) Multiply(args *Args, reply *int) error {
//			*reply = args.A * args.B
//			return nil
//		}
//
//		func (t *Arith) Divide(args *Args, quo *Quotient) error {
//			if args.B == 0 {
//				return errors.New("divide by zero")
//			}
//			quo.Quo = args.A / args.B
//			quo.Rem = args.A % args.B
//			return nil
//		}
//
// You can publish it in hprose like this:
//
//		service := rpc.NewHTTPService()
//		service.AddNetRPCMethods(new(Arith), rpc.Options{})
//		http.ListenAndServe(":8080", service)
//
//
// Then define stub struct in client:
//
//		type Stub struct {
//			// Synchronous call
//			Multiply func(args *Args) int
//			// Asynchronous call
//			Divide func(func(*Quotient, error), *Args)
//		}
//
// Then it can make a remote call:
//
//		client := rpc.NewClient("http://127.0.0.1:8080")
//		var stub *Stub
//		client.UseService(&stub)
//		fmt.Println(stub.Multiply(&Args{8, 7}))
//		stub.Divide(func(result *Quotient, err error) {
//			if err != nil {
//				log.Fatal("arith error:", err)
//			} else {
//				fmt.Println(result.Quo, result.Rem)
//			}
//		}, &Args{8, 7})
//
// You can also call it in other languages.
//
// Here is a simple example in JavaScript:
//
//		// The method name is case-insensitive
//		var methods = ['multiply', 'divide'];
//		var client = hprose.Client.create('http://127.0.0.1:8080/', methods);
//		// The first letter of the field name is convert to lowercase.
//      client.multiply({a:8, b:7}).then(function(product) {
//			return client.divide({a:product, b:6});
//		}).then(function(quotient) {
//			// The first letter of the field name is convert to lowercase.
//			console.log(quotient.quo, quotient.rem);
//		}).catch(function(err) {
//			// The returned err is here
//			console.error(err);
//		});
//
func (mm *methodManager) AddNetRPCMethods(rcvr interface{}, options Options) {
	if rcvr == nil {
		panic("rcvr can't be nil")
	}
	v := reflect.ValueOf(rcvr)
	t := v.Type()
	n := t.NumMethod()
	for i := 0; i < n; i++ {
		name := t.Method(i).Name
		method := v.Method(i)
		if method.CanInterface() {
			mm.addNetRPCMethod(name, method, options)
		}
	}
}

func (mm *methodManager) addNetRPCMethod(
	name string, method reflect.Value, options Options) {
	ft := method.Type()
	if ft.NumIn() != 2 || ft.IsVariadic() {
		panic("the method " + name + " must has two arguments")
	}
	if ft.In(1).Kind() != reflect.Ptr {
		panic("the second argument of method " + name + " must be a pointer")
	}
	if ft.NumOut() != 1 || ft.Out(0) != errorType {
		panic("the result type of method " + name + " must be error")
	}
	argsType := ft.In(0)
	resultType := ft.In(1).Elem()
	in := []reflect.Type{argsType}
	out := []reflect.Type{resultType, errorType}
	newft := reflect.FuncOf(in, out, false)
	newMethod := reflect.MakeFunc(newft, func(
		args []reflect.Value) (results []reflect.Value) {
		result := reflect.New(resultType)
		in := []reflect.Value{args[0], result}
		err := method.Call(in)[0]
		results = []reflect.Value{result.Elem(), err}
		return
	})
	mm.AddFunction(name, newMethod, options)
}

// Remove the published func or method by name
func (mm *methodManager) Remove(name string) {
	name = strings.ToLower(name)
	mm.Lock()
	n := len(mm.MethodNames)
	for i := 0; i < n; i++ {
		if strings.ToLower(mm.MethodNames[i]) == name {
			mm.MethodNames = append(mm.MethodNames[:i], mm.MethodNames[i+1:]...)
			break
		}
	}
	delete(mm.RemoteMethods, name)
	mm.Unlock()
}
<|endoftext|>


<commit_before>package eventexporter

import (
	"testing"

	"github.com/koding/eventexporter"
	"github.com/koding/eventexporter/eventexportertest"

	. "github.com/smartystreets/goconvey/convey"
)

func TestFakeExporter(t *testing.T) {
	Convey("When using FakeExporter", t, func() {
		event := &eventexporter.Event{Name: "test"}

		fakeExporter := eventexportertest.NewFakeExporter()
		err := fakeExporter.Send(event)

		Convey("Then it should save event for debugging", func() {
			So(err, ShouldBeNil)
			So(len(fakeExporter.Events), ShouldEqual, 1)
		})
	})
}
<commit_msg>fix import cycle<commit_after>package eventexporter_test

import (
	"testing"

	"github.com/koding/eventexporter/eventexportertest"

	. "github.com/koding/eventexporter"
	. "github.com/smartystreets/goconvey/convey"
)

func TestFakeExporter(t *testing.T) {
	Convey("When using FakeExporter", t, func() {
		event := &Event{Name: "test"}

		fakeExporter := eventexportertest.NewFakeExporter()
		err := fakeExporter.Send(event)

		Convey("Then it should save event for debugging", func() {
			So(err, ShouldBeNil)
			So(len(fakeExporter.Events), ShouldEqual, 1)
		})
	})
}
<|endoftext|>


<commit_before>/*
 * MinIO Go Library for Amazon S3 Compatible Cloud Storage
 * Copyright 2017 MinIO, Inc.
 *
 * 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 minio

import (
	"bufio"
	"context"
	"encoding/json"
	"net/http"
	"net/url"
	"time"

	"github.com/minio/minio-go/v6/pkg/s3utils"
)

// GetBucketNotification - get bucket notification at a given path.
func (c Client) GetBucketNotification(bucketName string) (bucketNotification BucketNotification, err error) {
	// Input validation.
	if err := s3utils.CheckValidBucketName(bucketName); err != nil {
		return BucketNotification{}, err
	}
	notification, err := c.getBucketNotification(bucketName)
	if err != nil {
		return BucketNotification{}, err
	}
	return notification, nil
}

// Request server for notification rules.
func (c Client) getBucketNotification(bucketName string) (BucketNotification, error) {
	urlValues := make(url.Values)
	urlValues.Set("notification", "")

	// Execute GET on bucket to list objects.
	resp, err := c.executeMethod(context.Background(), "GET", requestMetadata{
		bucketName:       bucketName,
		queryValues:      urlValues,
		contentSHA256Hex: emptySHA256Hex,
	})

	defer closeResponse(resp)
	if err != nil {
		return BucketNotification{}, err
	}
	return processBucketNotificationResponse(bucketName, resp)

}

// processes the GetNotification http response from the server.
func processBucketNotificationResponse(bucketName string, resp *http.Response) (BucketNotification, error) {
	if resp.StatusCode != http.StatusOK {
		errResponse := httpRespToErrorResponse(resp, bucketName, "")
		return BucketNotification{}, errResponse
	}
	var bucketNotification BucketNotification
	err := xmlDecoder(resp.Body, &bucketNotification)
	if err != nil {
		return BucketNotification{}, err
	}
	return bucketNotification, nil
}

// Indentity represents the user id, this is a compliance field.
type identity struct {
	PrincipalID string `json:"principalId"`
}

// Notification event bucket metadata.
type bucketMeta struct {
	Name          string   `json:"name"`
	OwnerIdentity identity `json:"ownerIdentity"`
	ARN           string   `json:"arn"`
}

// Notification event object metadata.
type objectMeta struct {
	Key       string `json:"key"`
	Size      int64  `json:"size,omitempty"`
	ETag      string `json:"eTag,omitempty"`
	VersionID string `json:"versionId,omitempty"`
	Sequencer string `json:"sequencer"`
}

// Notification event server specific metadata.
type eventMeta struct {
	SchemaVersion   string     `json:"s3SchemaVersion"`
	ConfigurationID string     `json:"configurationId"`
	Bucket          bucketMeta `json:"bucket"`
	Object          objectMeta `json:"object"`
}

// sourceInfo represents information on the client that
// triggered the event notification.
type sourceInfo struct {
	Host      string `json:"host"`
	Port      string `json:"port"`
	UserAgent string `json:"userAgent"`
}

// NotificationEvent represents an Amazon an S3 bucket notification event.
type NotificationEvent struct {
	EventVersion      string            `json:"eventVersion"`
	EventSource       string            `json:"eventSource"`
	AwsRegion         string            `json:"awsRegion"`
	EventTime         string            `json:"eventTime"`
	EventName         string            `json:"eventName"`
	UserIdentity      identity          `json:"userIdentity"`
	RequestParameters map[string]string `json:"requestParameters"`
	ResponseElements  map[string]string `json:"responseElements"`
	S3                eventMeta         `json:"s3"`
	Source            sourceInfo        `json:"source"`
}

// NotificationInfo - represents the collection of notification events, additionally
// also reports errors if any while listening on bucket notifications.
type NotificationInfo struct {
	Records []NotificationEvent
	Err     error
}

// ListenBucketNotification - listen on bucket notifications.
func (c Client) ListenBucketNotification(bucketName, prefix, suffix string, events []string, doneCh <-chan struct{}) <-chan NotificationInfo {
	notificationInfoCh := make(chan NotificationInfo, 1)
	// Only success, start a routine to start reading line by line.
	go func(notificationInfoCh chan<- NotificationInfo) {
		defer close(notificationInfoCh)

		// Validate the bucket name.
		if err := s3utils.CheckValidBucketName(bucketName); err != nil {
			notificationInfoCh <- NotificationInfo{
				Err: err,
			}
			return
		}

		// Check ARN partition to verify if listening bucket is supported
		if s3utils.IsAmazonEndpoint(*c.endpointURL) || s3utils.IsGoogleEndpoint(*c.endpointURL) {
			notificationInfoCh <- NotificationInfo{
				Err: ErrAPINotSupported("Listening for bucket notification is specific only to `minio` server endpoints"),
			}
			return
		}

		// Continuously run and listen on bucket notification.
		// Create a done channel to control 'ListObjects' go routine.
		retryDoneCh := make(chan struct{}, 1)

		// Indicate to our routine to exit cleanly upon return.
		defer close(retryDoneCh)

		// Prepare urlValues to pass into the request on every loop
		urlValues := make(url.Values)
		urlValues.Set("prefix", prefix)
		urlValues.Set("suffix", suffix)
		urlValues["events"] = events

		// Wait on the jitter retry loop.
		for range c.newRetryTimerContinous(time.Second, time.Second*30, MaxJitter, retryDoneCh) {
			// Execute GET on bucket to list objects.
			resp, err := c.executeMethod(context.Background(), "GET", requestMetadata{
				bucketName:       bucketName,
				queryValues:      urlValues,
				contentSHA256Hex: emptySHA256Hex,
			})
			if err != nil {
				notificationInfoCh <- NotificationInfo{
					Err: err,
				}
				return
			}

			// Validate http response, upon error return quickly.
			if resp.StatusCode != http.StatusOK {
				errResponse := httpRespToErrorResponse(resp, bucketName, "")
				notificationInfoCh <- NotificationInfo{
					Err: errResponse,
				}
				return
			}

			// Initialize a new bufio scanner, to read line by line.
			bio := bufio.NewScanner(resp.Body)

			// Unmarshal each line, returns marshalled values.
			for bio.Scan() {
				var notificationInfo NotificationInfo
				if err = json.Unmarshal(bio.Bytes(), &notificationInfo); err != nil {
					// Unexpected error during json unmarshal, send
					// the error to caller for actionable as needed.
					notificationInfoCh <- NotificationInfo{
						Err: err,
					}
					closeResponse(resp)
					continue
				}
				// Send notificationInfo
				select {
				case notificationInfoCh <- notificationInfo:
				case <-doneCh:
					closeResponse(resp)
					return
				}
			}
			if err = bio.Err(); err != nil {
				notificationInfoCh <- NotificationInfo{
					Err: err,
				}
			}
			// Close current connection before looping further.
			closeResponse(resp)
		}
	}(notificationInfoCh)

	// Returns the notification info channel, for caller to start reading from.
	return notificationInfoCh
}
<commit_msg>extract userMetadata from event response (#1229)<commit_after>/*
 * MinIO Go Library for Amazon S3 Compatible Cloud Storage
 * Copyright 2017-2020 MinIO, Inc.
 *
 * 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 minio

import (
	"bufio"
	"context"
	"encoding/json"
	"net/http"
	"net/url"
	"time"

	"github.com/minio/minio-go/v6/pkg/s3utils"
)

// GetBucketNotification - get bucket notification at a given path.
func (c Client) GetBucketNotification(bucketName string) (bucketNotification BucketNotification, err error) {
	// Input validation.
	if err := s3utils.CheckValidBucketName(bucketName); err != nil {
		return BucketNotification{}, err
	}
	notification, err := c.getBucketNotification(bucketName)
	if err != nil {
		return BucketNotification{}, err
	}
	return notification, nil
}

// Request server for notification rules.
func (c Client) getBucketNotification(bucketName string) (BucketNotification, error) {
	urlValues := make(url.Values)
	urlValues.Set("notification", "")

	// Execute GET on bucket to list objects.
	resp, err := c.executeMethod(context.Background(), "GET", requestMetadata{
		bucketName:       bucketName,
		queryValues:      urlValues,
		contentSHA256Hex: emptySHA256Hex,
	})

	defer closeResponse(resp)
	if err != nil {
		return BucketNotification{}, err
	}
	return processBucketNotificationResponse(bucketName, resp)

}

// processes the GetNotification http response from the server.
func processBucketNotificationResponse(bucketName string, resp *http.Response) (BucketNotification, error) {
	if resp.StatusCode != http.StatusOK {
		errResponse := httpRespToErrorResponse(resp, bucketName, "")
		return BucketNotification{}, errResponse
	}
	var bucketNotification BucketNotification
	err := xmlDecoder(resp.Body, &bucketNotification)
	if err != nil {
		return BucketNotification{}, err
	}
	return bucketNotification, nil
}

// Indentity represents the user id, this is a compliance field.
type identity struct {
	PrincipalID string `json:"principalId"`
}

// Notification event bucket metadata.
type bucketMeta struct {
	Name          string   `json:"name"`
	OwnerIdentity identity `json:"ownerIdentity"`
	ARN           string   `json:"arn"`
}

// Notification event object metadata.
type objectMeta struct {
	Key          string            `json:"key"`
	Size         int64             `json:"size,omitempty"`
	ETag         string            `json:"eTag,omitempty"`
	ContentType  string            `json:"contentType,omitempty"`
	UserMetadata map[string]string `json:"userMetadata,omitempty"`
	VersionID    string            `json:"versionId,omitempty"`
	Sequencer    string            `json:"sequencer"`
}

// Notification event server specific metadata.
type eventMeta struct {
	SchemaVersion   string     `json:"s3SchemaVersion"`
	ConfigurationID string     `json:"configurationId"`
	Bucket          bucketMeta `json:"bucket"`
	Object          objectMeta `json:"object"`
}

// sourceInfo represents information on the client that
// triggered the event notification.
type sourceInfo struct {
	Host      string `json:"host"`
	Port      string `json:"port"`
	UserAgent string `json:"userAgent"`
}

// NotificationEvent represents an Amazon an S3 bucket notification event.
type NotificationEvent struct {
	EventVersion      string            `json:"eventVersion"`
	EventSource       string            `json:"eventSource"`
	AwsRegion         string            `json:"awsRegion"`
	EventTime         string            `json:"eventTime"`
	EventName         string            `json:"eventName"`
	UserIdentity      identity          `json:"userIdentity"`
	RequestParameters map[string]string `json:"requestParameters"`
	ResponseElements  map[string]string `json:"responseElements"`
	S3                eventMeta         `json:"s3"`
	Source            sourceInfo        `json:"source"`
}

// NotificationInfo - represents the collection of notification events, additionally
// also reports errors if any while listening on bucket notifications.
type NotificationInfo struct {
	Records []NotificationEvent
	Err     error
}

// ListenBucketNotification - listen on bucket notifications.
func (c Client) ListenBucketNotification(bucketName, prefix, suffix string, events []string, doneCh <-chan struct{}) <-chan NotificationInfo {
	notificationInfoCh := make(chan NotificationInfo, 1)
	// Only success, start a routine to start reading line by line.
	go func(notificationInfoCh chan<- NotificationInfo) {
		defer close(notificationInfoCh)

		// Validate the bucket name.
		if err := s3utils.CheckValidBucketName(bucketName); err != nil {
			notificationInfoCh <- NotificationInfo{
				Err: err,
			}
			return
		}

		// Check ARN partition to verify if listening bucket is supported
		if s3utils.IsAmazonEndpoint(*c.endpointURL) || s3utils.IsGoogleEndpoint(*c.endpointURL) {
			notificationInfoCh <- NotificationInfo{
				Err: ErrAPINotSupported("Listening for bucket notification is specific only to `minio` server endpoints"),
			}
			return
		}

		// Continuously run and listen on bucket notification.
		// Create a done channel to control 'ListObjects' go routine.
		retryDoneCh := make(chan struct{}, 1)

		// Indicate to our routine to exit cleanly upon return.
		defer close(retryDoneCh)

		// Prepare urlValues to pass into the request on every loop
		urlValues := make(url.Values)
		urlValues.Set("prefix", prefix)
		urlValues.Set("suffix", suffix)
		urlValues["events"] = events

		// Wait on the jitter retry loop.
		for range c.newRetryTimerContinous(time.Second, time.Second*30, MaxJitter, retryDoneCh) {
			// Execute GET on bucket to list objects.
			resp, err := c.executeMethod(context.Background(), "GET", requestMetadata{
				bucketName:       bucketName,
				queryValues:      urlValues,
				contentSHA256Hex: emptySHA256Hex,
			})
			if err != nil {
				notificationInfoCh <- NotificationInfo{
					Err: err,
				}
				return
			}

			// Validate http response, upon error return quickly.
			if resp.StatusCode != http.StatusOK {
				errResponse := httpRespToErrorResponse(resp, bucketName, "")
				notificationInfoCh <- NotificationInfo{
					Err: errResponse,
				}
				return
			}

			// Initialize a new bufio scanner, to read line by line.
			bio := bufio.NewScanner(resp.Body)

			// Unmarshal each line, returns marshalled values.
			for bio.Scan() {
				var notificationInfo NotificationInfo
				if err = json.Unmarshal(bio.Bytes(), &notificationInfo); err != nil {
					// Unexpected error during json unmarshal, send
					// the error to caller for actionable as needed.
					notificationInfoCh <- NotificationInfo{
						Err: err,
					}
					closeResponse(resp)
					continue
				}
				// Send notificationInfo
				select {
				case notificationInfoCh <- notificationInfo:
				case <-doneCh:
					closeResponse(resp)
					return
				}
			}
			if err = bio.Err(); err != nil {
				notificationInfoCh <- NotificationInfo{
					Err: err,
				}
			}
			// Close current connection before looping further.
			closeResponse(resp)
		}
	}(notificationInfoCh)

	// Returns the notification info channel, for caller to start reading from.
	return notificationInfoCh
}
<|endoftext|>


<commit_before>package test

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/golang/protobuf/proto"

	"github.com/b2aio/typhon/server"
)

type StubServer struct {
	server.Server
	t *testing.T

	stubsMutex sync.RWMutex
	stubs      []*ServiceStub
}

type ServiceStub struct {
	ServiceName string
	Endpoint    string
	Handler     func(server.Request) (proto.Message, error)
}

// The stub server boots up a regular typhon server and registers a single
// endpoint that subscribes to every routing key
func NewStubServer(t *testing.T) *StubServer {

	stubServer := &StubServer{
		Server: server.NewAMQPServer(),
		t:      t,
	}

	// TODO `Name: "#"` is rather hokey
	stubServer.Init(&server.Config{Name: "#", Description: "Stub Server"})

	go stubServer.Run()

	select {
	case <-stubServer.NotifyConnected():
	case <-time.After(1 * time.Second):
		t.Fatalf("StubServer couldn't connect to RabbitMQ")
	}

	t.Log("[StubServer] Connected to RabbitMQ")

	stubServer.RegisterEndpoint(&server.Endpoint{
		Name: ".*", // TODO Name is not well-named
		Handler: func(req server.Request) (proto.Message, error) {
			return stubServer.handleRequest(req)
		},
	})

	return stubServer
}

// StubResponse is a convenience method to quickly set up stubs that return a fixed value
func (stubServer *StubServer) StubResponse(serviceName, endpoint string, returnValue proto.Message) {
	stubServer.stubResponseAndError(serviceName, endpoint, returnValue, nil)
}

// StubError is a convenience method to stub out a service error
func (stubServer *StubServer) StubError(serviceName, endpoint string, err error) {
	stubServer.stubResponseAndError(serviceName, endpoint, nil, err)
}

// stubResponseAndError registers a stub that returns the passed response and error
func (stubServer *StubServer) stubResponseAndError(serviceName, endpoint string, returnValue proto.Message, err error) {
	stubServer.RegisterStub(&ServiceStub{
		ServiceName: serviceName,
		Endpoint:    endpoint,
		Handler: func(_ server.Request) (proto.Message, error) {
			return returnValue, err
		},
	})
}

// Registers a stub with the server
func (stubServer *StubServer) RegisterStub(stub *ServiceStub) {
	stubServer.stubsMutex.Lock()
	stubServer.t.Logf("[StubServer] Registered stub for %s.%s", stub.ServiceName, stub.Endpoint)
	defer stubServer.stubsMutex.Unlock()
	stubServer.stubs = append(stubServer.stubs, stub)
}

// Clear out all server stubs. Test suites should run this between tests
func (stubServer *StubServer) ResetStubs() {
	stubServer.stubsMutex.Lock()
	defer stubServer.stubsMutex.Unlock()
	stubServer.stubs = nil
	stubServer.t.Log("[StubServer] Stubs cleared")
}

// Finds the relevant endpoint stub (if any), and calls its handler function
func (stubServer *StubServer) handleRequest(req server.Request) (proto.Message, error) {

	stubServer.t.Logf("[StubServer] Handling request for %s", req.Service(), req.Endpoint())
	stubServer.stubsMutex.RLock()
	defer stubServer.stubsMutex.RUnlock()

	// determine which endpoint to use
	for _, stub := range stubServer.stubs {
		if stub.ServiceName == req.Service() && stub.Endpoint == req.Endpoint() {
			return stub.Handler(req)
		}
	}
	return nil, fmt.Errorf("No stub found for routing service name %s and endpoint %s", req.Service(), req.Endpoint())
}
<commit_msg>Stub server: Add comments and update to match conventions<commit_after>package test

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/golang/protobuf/proto"

	"github.com/b2aio/typhon/server"
)

type StubServer struct {
	server.Server
	t *testing.T

	stubsMutex sync.RWMutex
	stubs      []*ServiceStub
}

type ServiceStub struct {
	ServiceName string
	Endpoint    string
	Handler     func(server.Request) (proto.Message, error)
}

// NewStubServer boots up a regular typhon server and registers a single
// endpoint that subscribes to every routing key
func NewStubServer(t *testing.T) *StubServer {

	s := &StubServer{
		Server: server.NewAMQPServer(),
		t:      t,
	}

	// TODO `Name: "#"` is rather hokey
	s.Init(&server.Config{Name: "#", Description: "Stub Server"})

	go s.Run()

	select {
	case <-s.NotifyConnected():
	case <-time.After(1 * time.Second):
		t.Fatalf("StubServer couldn't connect to RabbitMQ")
	}

	t.Log("[StubServer] Connected to RabbitMQ")

	s.RegisterEndpoint(&server.Endpoint{
		Name: ".*", // TODO Name is not well-named
		Handler: func(req server.Request) (proto.Message, error) {
			return s.handleRequest(req)
		},
	})

	return s
}

// StubResponse is a convenience method to quickly set up stubs that return a fixed value
func (s *StubServer) StubResponse(serviceName, endpoint string, returnValue proto.Message) {
	s.stubResponseAndError(serviceName, endpoint, returnValue, nil)
}

// StubError is a convenience method to stub out a service error
func (s *StubServer) StubError(serviceName, endpoint string, err error) {
	s.stubResponseAndError(serviceName, endpoint, nil, err)
}

// stubResponseAndError registers a stub that returns the passed response and error
func (s *StubServer) stubResponseAndError(serviceName, endpoint string, returnValue proto.Message, err error) {
	s.RegisterStub(&ServiceStub{
		ServiceName: serviceName,
		Endpoint:    endpoint,
		Handler: func(_ server.Request) (proto.Message, error) {
			return returnValue, err
		},
	})
}

// RegisterStub with the server
func (s *StubServer) RegisterStub(stub *ServiceStub) {
	s.stubsMutex.Lock()
	s.t.Logf("[StubServer] Registered stub for %s.%s", stub.ServiceName, stub.Endpoint)
	defer s.stubsMutex.Unlock()
	s.stubs = append(s.stubs, stub)
}

// ResetStubs clears out all server stubs. Test suites should run this between tests
func (s *StubServer) ResetStubs() {
	s.stubsMutex.Lock()
	defer s.stubsMutex.Unlock()
	s.stubs = nil
	s.t.Log("[StubServer] Stubs cleared")
}

// Close cleanly shuts down the stub server
func (s *StubServer) Close() {
	s.Server.Close()
}

// Finds the relevant endpoint stub (if any), and calls its handler function
func (s *StubServer) handleRequest(req server.Request) (proto.Message, error) {

	s.t.Logf("[StubServer] Handling request for %s.%s", req.Service(), req.Endpoint())
	s.stubsMutex.RLock()
	defer s.stubsMutex.RUnlock()

	// determine which endpoint to use
	for _, stub := range s.stubs {
		if stub.ServiceName == req.Service() && stub.Endpoint == req.Endpoint() {
			return stub.Handler(req)
		}
	}
	return nil, fmt.Errorf("No stub found for routing service name %s and endpoint %s", req.Service(), req.Endpoint())
}
<|endoftext|>