<commit_before>package implements

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"io/ioutil"
	"path"
	"regexp"
	"strings"
)

// ignoreCCalls is a map of calls from "C" that we know we can just ignore.
// ignoring this stuff makes the number of tracked entries smaller, fewer logs,
// and less stuff to iterate over in other code.
var ignoreCCalls = map[string]bool{
	// some special cgo calls
	"CString":  true,
	"CBytes":   true,
	"GoString": true,
	"GoBytes":  true,
	// common utility functions
	"free": true,
	// common types
	"int":       true,
	"int64_t":   true,
	"uint64_t":  true,
	"size_t":    true,
	"ssize_t":   true,
	"uintptr_t": true,
}

type goFunction struct {
	shortName       string
	fullName        string
	comment         string
	implementsCFunc string
	callsCFunc      string
	isDeprecated    bool
	isPreview       bool

	endPos token.Pos
}

type visitor struct {
	currentFunc *goFunction

	callMap    map[string]*goFunction
	docMap     map[string]*goFunction
	deprecated []*goFunction
	preview    []*goFunction
	stable     []*goFunction
}

func newVisitor() *visitor {
	return &visitor{
		callMap:    map[string]*goFunction{},
		docMap:     map[string]*goFunction{},
		deprecated: []*goFunction{},
		preview:    []*goFunction{},
		stable:     []*goFunction{},
	}
}

func funcDeclFullName(fdec *ast.FuncDecl) string {
	if fdec.Recv == nil {
		return fdec.Name.Name
	}
	if len(fdec.Recv.List) != 1 {
		return fdec.Name.Name
	}
	typeName := "UNKNOWN!"
	switch t := fdec.Recv.List[0].Type.(type) {
	case *ast.StarExpr:
		typeName = t.X.(*ast.Ident).Name
	case *ast.Ident:
		typeName = t.Name
	}
	return fmt.Sprintf("%s.%s", typeName, fdec.Name.Name)
}

func readDocComment(fdec *ast.FuncDecl, gfunc *goFunction) {
	gfunc.comment = fdec.Doc.Text()
	lines := strings.Split(gfunc.comment, "\n")
	for i := range lines {
		if strings.Contains(lines[i], "DEPRECATED") {
			gfunc.isDeprecated = true
			logger.Printf("marked deprecated: %s\n", fdec.Name.Name)
		}
		if strings.Contains(lines[i], "PREVIEW") {
			gfunc.isPreview = true
			logger.Printf("marked preview: %s\n", fdec.Name.Name)
		}

		if lines[i] == "Implements:" {
			cfunc := cfuncFromComment(lines[i+1])
			if cfunc == "" {
				return
			}
			gfunc.implementsCFunc = cfunc
			logger.Printf("implements c function %s: %s\n", fdec.Name.Name, cfunc)
		}
	}
}

func isPublic(gfunc *goFunction) bool {
	return ast.IsExported(gfunc.shortName)
}

func (v *visitor) checkCalled(s *ast.SelectorExpr) {
	ident, ok := s.X.(*ast.Ident)
	if !ok {
		return
	}
	if "C" == ident.String() {
		cname := s.Sel.String()
		if ignoreCCalls[cname] {
			return
		}
		v.callMap[cname] = v.currentFunc
		logger.Printf("updated %s in call map\n", s.Sel.String())
	}
}

func (v *visitor) Visit(node ast.Node) ast.Visitor {
	switch {
	case node == nil:
		return nil
	case v.currentFunc == nil:
	case node.Pos() > v.currentFunc.endPos:
		logger.Printf("left function %v\n", v.currentFunc.shortName)
		v.currentFunc = nil
	}

	switch n := node.(type) {
	case *ast.File:
		v.currentFunc = nil
		return v
	case *ast.FuncDecl:
		logger.Printf("checking function: %v\n", n.Name.Name)
		gfunc := &goFunction{
			shortName: n.Name.Name,
			fullName:  funcDeclFullName(n),
			endPos:    n.End(),
		}
		readDocComment(n, gfunc)
		v.currentFunc = gfunc
		if isPublic(gfunc) {
			if gfunc.isDeprecated {
				v.deprecated = append(v.deprecated, gfunc)
			} else if gfunc.isPreview {
				v.preview = append(v.preview, gfunc)
			} else {
				v.stable = append(v.stable, gfunc)
			}
		}
		return v
	case *ast.CallExpr:
		if v.currentFunc == nil {
			return nil
		}
		if s, ok := n.Fun.(*ast.SelectorExpr); ok {
			v.checkCalled(s)
		}
	}
	if v.currentFunc != nil {
		return v
	}
	return nil
}

func cfuncFromComment(ctext string) string {
	m := regexp.MustCompile(` ([a-zA-Z0-9_]+)\(`).FindAllSubmatch([]byte(ctext), 1)
	if len(m) < 1 {
		return ""
	}
	return string(m[0][1])
}

// CephGoFunctions will look for C functions called by the code code and
// update the found functions for the package within the inspector.
func CephGoFunctions(source, packageName string, ii *Inspector) error {
	p, err := build.Import("./"+packageName, source, 0)
	if err != nil {
		return err
	}

	toCheck := []string{}
	toCheck = append(toCheck, p.GoFiles...)
	toCheck = append(toCheck, p.CgoFiles...)
	for _, fname := range toCheck {
		logger.Printf("Reading go file: %v\n", fname)
		src, err := ioutil.ReadFile(path.Join(p.Dir, fname))
		if err != nil {
			return err
		}
		fset := token.NewFileSet()
		f, err := parser.ParseFile(
			fset,
			fname,
			src,
			parser.ParseComments|parser.AllErrors)
		if err != nil {
			return err
		}
		ast.Walk(ii.visitor, f)
	}
	return nil
}
<commit_msg>implements: always include the ceph_preview build tag<commit_after>package implements

import (
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"io/ioutil"
	"path"
	"regexp"
	"strings"
)

// ignoreCCalls is a map of calls from "C" that we know we can just ignore.
// ignoring this stuff makes the number of tracked entries smaller, fewer logs,
// and less stuff to iterate over in other code.
var ignoreCCalls = map[string]bool{
	// some special cgo calls
	"CString":  true,
	"CBytes":   true,
	"GoString": true,
	"GoBytes":  true,
	// common utility functions
	"free": true,
	// common types
	"int":       true,
	"int64_t":   true,
	"uint64_t":  true,
	"size_t":    true,
	"ssize_t":   true,
	"uintptr_t": true,
}

type goFunction struct {
	shortName       string
	fullName        string
	comment         string
	implementsCFunc string
	callsCFunc      string
	isDeprecated    bool
	isPreview       bool

	endPos token.Pos
}

type visitor struct {
	currentFunc *goFunction

	callMap    map[string]*goFunction
	docMap     map[string]*goFunction
	deprecated []*goFunction
	preview    []*goFunction
	stable     []*goFunction
}

func newVisitor() *visitor {
	return &visitor{
		callMap:    map[string]*goFunction{},
		docMap:     map[string]*goFunction{},
		deprecated: []*goFunction{},
		preview:    []*goFunction{},
		stable:     []*goFunction{},
	}
}

func funcDeclFullName(fdec *ast.FuncDecl) string {
	if fdec.Recv == nil {
		return fdec.Name.Name
	}
	if len(fdec.Recv.List) != 1 {
		return fdec.Name.Name
	}
	typeName := "UNKNOWN!"
	switch t := fdec.Recv.List[0].Type.(type) {
	case *ast.StarExpr:
		typeName = t.X.(*ast.Ident).Name
	case *ast.Ident:
		typeName = t.Name
	}
	return fmt.Sprintf("%s.%s", typeName, fdec.Name.Name)
}

func readDocComment(fdec *ast.FuncDecl, gfunc *goFunction) {
	gfunc.comment = fdec.Doc.Text()
	lines := strings.Split(gfunc.comment, "\n")
	for i := range lines {
		if strings.Contains(lines[i], "DEPRECATED") {
			gfunc.isDeprecated = true
			logger.Printf("marked deprecated: %s\n", fdec.Name.Name)
		}
		if strings.Contains(lines[i], "PREVIEW") {
			gfunc.isPreview = true
			logger.Printf("marked preview: %s\n", fdec.Name.Name)
		}

		if lines[i] == "Implements:" {
			cfunc := cfuncFromComment(lines[i+1])
			if cfunc == "" {
				return
			}
			gfunc.implementsCFunc = cfunc
			logger.Printf("implements c function %s: %s\n", fdec.Name.Name, cfunc)
		}
	}
}

func isPublic(gfunc *goFunction) bool {
	return ast.IsExported(gfunc.shortName)
}

func (v *visitor) checkCalled(s *ast.SelectorExpr) {
	ident, ok := s.X.(*ast.Ident)
	if !ok {
		return
	}
	if "C" == ident.String() {
		cname := s.Sel.String()
		if ignoreCCalls[cname] {
			return
		}
		v.callMap[cname] = v.currentFunc
		logger.Printf("updated %s in call map\n", s.Sel.String())
	}
}

func (v *visitor) Visit(node ast.Node) ast.Visitor {
	switch {
	case node == nil:
		return nil
	case v.currentFunc == nil:
	case node.Pos() > v.currentFunc.endPos:
		logger.Printf("left function %v\n", v.currentFunc.shortName)
		v.currentFunc = nil
	}

	switch n := node.(type) {
	case *ast.File:
		v.currentFunc = nil
		return v
	case *ast.FuncDecl:
		logger.Printf("checking function: %v\n", n.Name.Name)
		gfunc := &goFunction{
			shortName: n.Name.Name,
			fullName:  funcDeclFullName(n),
			endPos:    n.End(),
		}
		readDocComment(n, gfunc)
		v.currentFunc = gfunc
		if isPublic(gfunc) {
			if gfunc.isDeprecated {
				v.deprecated = append(v.deprecated, gfunc)
			} else if gfunc.isPreview {
				v.preview = append(v.preview, gfunc)
			} else {
				v.stable = append(v.stable, gfunc)
			}
		}
		return v
	case *ast.CallExpr:
		if v.currentFunc == nil {
			return nil
		}
		if s, ok := n.Fun.(*ast.SelectorExpr); ok {
			v.checkCalled(s)
		}
	}
	if v.currentFunc != nil {
		return v
	}
	return nil
}

func cfuncFromComment(ctext string) string {
	m := regexp.MustCompile(` ([a-zA-Z0-9_]+)\(`).FindAllSubmatch([]byte(ctext), 1)
	if len(m) < 1 {
		return ""
	}
	return string(m[0][1])
}

// CephGoFunctions will look for C functions called by the code code and
// update the found functions for the package within the inspector.
func CephGoFunctions(source, packageName string, ii *Inspector) error {
	p, err := build.Import("./"+packageName, source, 0)
	if err != nil {
		return err
	}

	toCheck := []string{}
	toCheck = append(toCheck, p.GoFiles...)
	toCheck = append(toCheck, p.CgoFiles...)
	for _, fname := range p.IgnoredGoFiles {
		logger.Printf("Ignored go file: %v\n", fname)
	}
	for _, fname := range toCheck {
		logger.Printf("Reading go file: %v\n", fname)
		src, err := ioutil.ReadFile(path.Join(p.Dir, fname))
		if err != nil {
			return err
		}
		fset := token.NewFileSet()
		f, err := parser.ParseFile(
			fset,
			fname,
			src,
			parser.ParseComments|parser.AllErrors)
		if err != nil {
			return err
		}
		ast.Walk(ii.visitor, f)
	}
	return nil
}

func init() {
	build.Default.BuildTags = append(build.Default.BuildTags, "ceph_preview")
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Audrius Butkevicius and Contributors (see the CONTRIBUTORS file).

package client

import (
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/url"
	"time"

	syncthingprotocol "github.com/syncthing/protocol"
	"github.com/syncthing/relaysrv/protocol"
	"github.com/syncthing/syncthing/lib/sync"
)

type ProtocolClient struct {
	URI         *url.URL
	Invitations chan protocol.SessionInvitation

	closeInvitationsOnFinish bool

	config *tls.Config

	timeout time.Duration

	stop    chan struct{}
	stopped chan struct{}

	conn *tls.Conn

	mut       sync.RWMutex
	connected bool
}

func NewProtocolClient(uri *url.URL, certs []tls.Certificate, invitations chan protocol.SessionInvitation) *ProtocolClient {
	closeInvitationsOnFinish := false
	if invitations == nil {
		closeInvitationsOnFinish = true
		invitations = make(chan protocol.SessionInvitation)
	}

	return &ProtocolClient{
		URI:         uri,
		Invitations: invitations,

		closeInvitationsOnFinish: closeInvitationsOnFinish,

		config: configForCerts(certs),

		timeout: time.Minute * 2,

		stop:    make(chan struct{}),
		stopped: make(chan struct{}),

		mut:       sync.NewRWMutex(),
		connected: false,
	}
}

func (c *ProtocolClient) Serve() {
	c.stop = make(chan struct{})
	c.stopped = make(chan struct{})
	defer close(c.stopped)

	if err := c.connect(); err != nil {
		if debug {
			l.Debugln("Relay connect:", err)
		}
		return
	}

	if debug {
		l.Debugln(c, "connected", c.conn.RemoteAddr())
	}

	if err := c.join(); err != nil {
		c.conn.Close()
		l.Infoln("Relay join:", err)
		return
	}

	if err := c.conn.SetDeadline(time.Time{}); err != nil {
		l.Infoln("Relay set deadline:", err)
		return
	}

	if debug {
		l.Debugln(c, "joined", c.conn.RemoteAddr(), "via", c.conn.LocalAddr())
	}

	defer c.cleanup()
	c.mut.Lock()
	c.connected = true
	c.mut.Unlock()

	messages := make(chan interface{})
	errors := make(chan error, 1)

	go messageReader(c.conn, messages, errors)

	timeout := time.NewTimer(c.timeout)

	for {
		select {
		case message := <-messages:
			timeout.Reset(c.timeout)
			if debug {
				log.Printf("%s received message %T", c, message)
			}

			switch msg := message.(type) {
			case protocol.Ping:
				if err := protocol.WriteMessage(c.conn, protocol.Pong{}); err != nil {
					l.Infoln("Relay write:", err)
					return

				}
				if debug {
					l.Debugln(c, "sent pong")
				}

			case protocol.SessionInvitation:
				ip := net.IP(msg.Address)
				if len(ip) == 0 || ip.IsUnspecified() {
					msg.Address = c.conn.RemoteAddr().(*net.TCPAddr).IP[:]
				}
				c.Invitations <- msg

			default:
				l.Infoln("Relay: protocol error: unexpected message %v", msg)
				return
			}

		case <-c.stop:
			if debug {
				l.Debugln(c, "stopping")
			}
			return

		case err := <-errors:
			l.Infoln("Relay received:", err)
			return

		case <-timeout.C:
			if debug {
				l.Debugln(c, "timed out")
			}
			return
		}
	}
}

func (c *ProtocolClient) Stop() {
	if c.stop == nil {
		return
	}

	close(c.stop)
	<-c.stopped
}

func (c *ProtocolClient) StatusOK() bool {
	c.mut.RLock()
	con := c.connected
	c.mut.RUnlock()
	return con
}

func (c *ProtocolClient) String() string {
	return fmt.Sprintf("ProtocolClient@%p", c)
}

func (c *ProtocolClient) connect() error {
	if c.URI.Scheme != "relay" {
		return fmt.Errorf("Unsupported relay schema:", c.URI.Scheme)
	}

	conn, err := tls.Dial("tcp", c.URI.Host, c.config)
	if err != nil {
		return err
	}

	if err := conn.SetDeadline(time.Now().Add(10 * time.Second)); err != nil {
		conn.Close()
		return err
	}

	if err := performHandshakeAndValidation(conn, c.URI); err != nil {
		conn.Close()
		return err
	}

	c.conn = conn
	return nil
}

func (c *ProtocolClient) cleanup() {
	if c.closeInvitationsOnFinish {
		close(c.Invitations)
		c.Invitations = make(chan protocol.SessionInvitation)
	}

	if debug {
		l.Debugln(c, "cleaning up")
	}

	c.mut.Lock()
	c.connected = false
	c.mut.Unlock()

	c.conn.Close()
}

func (c *ProtocolClient) join() error {
	if err := protocol.WriteMessage(c.conn, protocol.JoinRelayRequest{}); err != nil {
		return err
	}

	message, err := protocol.ReadMessage(c.conn)
	if err != nil {
		return err
	}

	switch msg := message.(type) {
	case protocol.Response:
		if msg.Code != 0 {
			return fmt.Errorf("Incorrect response code %d: %s", msg.Code, msg.Message)
		}

	default:
		return fmt.Errorf("protocol error: expecting response got %v", msg)
	}

	return nil
}

func performHandshakeAndValidation(conn *tls.Conn, uri *url.URL) error {
	if err := conn.Handshake(); err != nil {
		return err
	}

	cs := conn.ConnectionState()
	if !cs.NegotiatedProtocolIsMutual || cs.NegotiatedProtocol != protocol.ProtocolName {
		return fmt.Errorf("protocol negotiation error")
	}

	q := uri.Query()
	relayIDs := q.Get("id")
	if relayIDs != "" {
		relayID, err := syncthingprotocol.DeviceIDFromString(relayIDs)
		if err != nil {
			return fmt.Errorf("relay address contains invalid verification id: %s", err)
		}

		certs := cs.PeerCertificates
		if cl := len(certs); cl != 1 {
			return fmt.Errorf("unexpected certificate count: %d", cl)
		}

		remoteID := syncthingprotocol.NewDeviceID(certs[0].Raw)
		if remoteID != relayID {
			return fmt.Errorf("relay id does not match. Expected %v got %v", relayID, remoteID)
		}
	}

	return nil
}

func messageReader(conn net.Conn, messages chan<- interface{}, errors chan<- error) {
	for {
		msg, err := protocol.ReadMessage(conn)
		if err != nil {
			errors <- err
			return
		}
		messages <- msg
	}
}
<commit_msg>Connected clients should know their own latency<commit_after>// Copyright (C) 2015 Audrius Butkevicius and Contributors (see the CONTRIBUTORS file).

package client

import (
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/url"
	"time"

	syncthingprotocol "github.com/syncthing/protocol"
	"github.com/syncthing/relaysrv/protocol"
	"github.com/syncthing/syncthing/lib/sync"
)

type ProtocolClient struct {
	URI         *url.URL
	Invitations chan protocol.SessionInvitation

	closeInvitationsOnFinish bool

	config *tls.Config

	timeout time.Duration

	stop    chan struct{}
	stopped chan struct{}

	conn *tls.Conn

	mut       sync.RWMutex
	connected bool
	latency   time.Duration
}

func NewProtocolClient(uri *url.URL, certs []tls.Certificate, invitations chan protocol.SessionInvitation) *ProtocolClient {
	closeInvitationsOnFinish := false
	if invitations == nil {
		closeInvitationsOnFinish = true
		invitations = make(chan protocol.SessionInvitation)
	}

	return &ProtocolClient{
		URI:         uri,
		Invitations: invitations,

		closeInvitationsOnFinish: closeInvitationsOnFinish,

		config: configForCerts(certs),

		timeout: time.Minute * 2,

		stop:    make(chan struct{}),
		stopped: make(chan struct{}),

		mut:       sync.NewRWMutex(),
		connected: false,
	}
}

func (c *ProtocolClient) Serve() {
	c.stop = make(chan struct{})
	c.stopped = make(chan struct{})
	defer close(c.stopped)

	if err := c.connect(); err != nil {
		if debug {
			l.Debugln("Relay connect:", err)
		}
		return
	}

	if debug {
		l.Debugln(c, "connected", c.conn.RemoteAddr())
	}

	if err := c.join(); err != nil {
		c.conn.Close()
		l.Infoln("Relay join:", err)
		return
	}

	if err := c.conn.SetDeadline(time.Time{}); err != nil {
		l.Infoln("Relay set deadline:", err)
		return
	}

	if debug {
		l.Debugln(c, "joined", c.conn.RemoteAddr(), "via", c.conn.LocalAddr())
	}

	defer c.cleanup()
	c.mut.Lock()
	c.connected = true
	c.mut.Unlock()

	messages := make(chan interface{})
	errors := make(chan error, 1)

	go messageReader(c.conn, messages, errors)

	timeout := time.NewTimer(c.timeout)

	for {
		select {
		case message := <-messages:
			timeout.Reset(c.timeout)
			if debug {
				log.Printf("%s received message %T", c, message)
			}

			switch msg := message.(type) {
			case protocol.Ping:
				if err := protocol.WriteMessage(c.conn, protocol.Pong{}); err != nil {
					l.Infoln("Relay write:", err)
					return

				}
				if debug {
					l.Debugln(c, "sent pong")
				}

			case protocol.SessionInvitation:
				ip := net.IP(msg.Address)
				if len(ip) == 0 || ip.IsUnspecified() {
					msg.Address = c.conn.RemoteAddr().(*net.TCPAddr).IP[:]
				}
				c.Invitations <- msg

			default:
				l.Infoln("Relay: protocol error: unexpected message %v", msg)
				return
			}

		case <-c.stop:
			if debug {
				l.Debugln(c, "stopping")
			}
			return

		case err := <-errors:
			l.Infoln("Relay received:", err)
			return

		case <-timeout.C:
			if debug {
				l.Debugln(c, "timed out")
			}
			return
		}
	}
}

func (c *ProtocolClient) Stop() {
	if c.stop == nil {
		return
	}

	close(c.stop)
	<-c.stopped
}

func (c *ProtocolClient) StatusOK() bool {
	c.mut.RLock()
	con := c.connected
	c.mut.RUnlock()
	return con
}

func (c *ProtocolClient) Latency() time.Duration {
	c.mut.RLock()
	lat := c.latency
	c.mut.RUnlock()
	return lat
}

func (c *ProtocolClient) String() string {
	return fmt.Sprintf("ProtocolClient@%p", c)
}

func (c *ProtocolClient) connect() error {
	if c.URI.Scheme != "relay" {
		return fmt.Errorf("Unsupported relay schema:", c.URI.Scheme)
	}

	t0 := time.Now()
	tcpConn, err := net.Dial("tcp", c.URI.Host)
	if err != nil {
		return err
	}

	c.mut.Lock()
	c.latency = time.Since(t0)
	c.mut.Unlock()

	conn := tls.Client(tcpConn, c.config)
	if err = conn.Handshake(); err != nil {
		return err
	}

	if err := conn.SetDeadline(time.Now().Add(10 * time.Second)); err != nil {
		conn.Close()
		return err
	}

	if err := performHandshakeAndValidation(conn, c.URI); err != nil {
		conn.Close()
		return err
	}

	c.conn = conn
	return nil
}

func (c *ProtocolClient) cleanup() {
	if c.closeInvitationsOnFinish {
		close(c.Invitations)
		c.Invitations = make(chan protocol.SessionInvitation)
	}

	if debug {
		l.Debugln(c, "cleaning up")
	}

	c.mut.Lock()
	c.connected = false
	c.mut.Unlock()

	c.conn.Close()
}

func (c *ProtocolClient) join() error {
	if err := protocol.WriteMessage(c.conn, protocol.JoinRelayRequest{}); err != nil {
		return err
	}

	message, err := protocol.ReadMessage(c.conn)
	if err != nil {
		return err
	}

	switch msg := message.(type) {
	case protocol.Response:
		if msg.Code != 0 {
			return fmt.Errorf("Incorrect response code %d: %s", msg.Code, msg.Message)
		}

	default:
		return fmt.Errorf("protocol error: expecting response got %v", msg)
	}

	return nil
}

func performHandshakeAndValidation(conn *tls.Conn, uri *url.URL) error {
	if err := conn.Handshake(); err != nil {
		return err
	}

	cs := conn.ConnectionState()
	if !cs.NegotiatedProtocolIsMutual || cs.NegotiatedProtocol != protocol.ProtocolName {
		return fmt.Errorf("protocol negotiation error")
	}

	q := uri.Query()
	relayIDs := q.Get("id")
	if relayIDs != "" {
		relayID, err := syncthingprotocol.DeviceIDFromString(relayIDs)
		if err != nil {
			return fmt.Errorf("relay address contains invalid verification id: %s", err)
		}

		certs := cs.PeerCertificates
		if cl := len(certs); cl != 1 {
			return fmt.Errorf("unexpected certificate count: %d", cl)
		}

		remoteID := syncthingprotocol.NewDeviceID(certs[0].Raw)
		if remoteID != relayID {
			return fmt.Errorf("relay id does not match. Expected %v got %v", relayID, remoteID)
		}
	}

	return nil
}

func messageReader(conn net.Conn, messages chan<- interface{}, errors chan<- error) {
	for {
		msg, err := protocol.ReadMessage(conn)
		if err != nil {
			errors <- err
			return
		}
		messages <- msg
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"gopkg.in/natefinch/lumberjack.v2"
	"k8s.io/klog"
)

func main() {
	terminationLog := flag.String("termination-log-file", "", "Write logs after SIGTERM to this file (in addition to stderr)")
	terminationLock := flag.String("termination-touch-file", "", "Touch this file on SIGTERM and delete on termination")

	klog.InitFlags(nil)
	flag.Set("v", "9")

	// never log to stderr, only through our termination log writer (which sends it also to stderr)
	flag.Set("logtostderr", "false")
	flag.Set("stderrthreshold", "99")

	flag.Parse()
	args := flag.CommandLine.Args()

	if len(args) == 0 {
		fmt.Println("Missing command line")
		os.Exit(1)
	}

	// use special tee-like writer when termination log is set
	termCh := make(chan struct{})
	var stderr io.Writer = os.Stderr
	if len(*terminationLog) > 0 {
		stderr = &terminationFileWriter{
			Writer:             os.Stderr,
			fn:                 *terminationLog,
			startFileLoggingCh: termCh,
		}

		// do the klog file writer dance: klog writes to all outputs of lower
		// severity. No idea why. So we discard for anything other than info.
		// Otherwise, we would see errors multiple times.
		klog.SetOutput(ioutil.Discard)
		klog.SetOutputBySeverity("INFO", stderr)
	}

	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = stderr

	// forward SIGTERM and SIGINT to child
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		for s := range sigCh {
			select {
			case <-termCh:
			default:
				close(termCh)
			}

			klog.Infof("Received signal %s. Forwarding to sub-process %q.", s, args[0])

			if len(*terminationLock) > 0 {
				klog.Infof("Touching termination lock file %q", *terminationLock)
				if err := touch(*terminationLock); err != nil {
					klog.Infof("error touching %s: %v", *terminationLock, err)
					// keep going
				}
			}

			cmd.Process.Signal(s)
		}
	}()

	klog.Infof("Launching sub-process %q", cmd)
	rc := 0
	if err := cmd.Run(); err != nil {
		if exitError, ok := err.(*exec.ExitError); ok {
			rc = exitError.ExitCode()
		} else {
			klog.Infof("Failed to launch %s: %v", args[0], err)
			os.Exit(255)
		}
	}

	// remove signal handling
	signal.Stop(sigCh)
	close(sigCh)
	wg.Wait()

	if len(*terminationLock) > 0 {
		klog.Infof("Deleting termination lock file %q", *terminationLock)
		os.Remove(*terminationLock)
	}

	klog.Infof("Termination finished with exit code %d", rc)
	os.Exit(rc)
}

// terminationFileWriter forwards everything to the embedded writer. When
// startFileLoggingCh is closed, everything is appended to the given file name
// in addition.
type terminationFileWriter struct {
	io.Writer
	fn                 string
	startFileLoggingCh <-chan struct{}

	logger io.Writer
}

func (w *terminationFileWriter) Write(bs []byte) (int, error) {
	// temporary hack to avoid logging sensitive tokens.
	// TODO: drop when we moved to a non-sensitive storage format
	if strings.Contains(string(bs), "URI=\"/apis/oauth.openshift.io/v1/oauthaccesstokens/") || strings.Contains(string(bs), "URI=\"/apis/oauth.openshift.io/v1/oauthauthorizetokens/") {
		return len(bs), nil
	}

	select {
	case <-w.startFileLoggingCh:
		if w.logger == nil {
			l := &lumberjack.Logger{
				Filename:   w.fn,
				MaxSize:    100,
				MaxBackups: 3,
				MaxAge:     28,
				Compress:   false,
			}
			w.logger = l
			fmt.Fprintf(os.Stderr, "Copying termination logs to %q\n", w.fn)
		}
		if n, err := w.logger.Write(bs); err != nil {
			return n, err
		} else if n != len(bs) {
			return n, io.ErrShortWrite
		}
	default:
	}

	return w.Writer.Write(bs)
}

func touch(fn string) error {
	_, err := os.Stat(fn)
	if os.IsNotExist(err) {
		file, err := os.Create(fn)
		if err != nil {
			return err
		}
		defer file.Close()
		return nil
	}

	currentTime := time.Now().Local()
	return os.Chtimes(fn, currentTime, currentTime)
}
<commit_msg>watch-termination: create lock file on launch to catch every hard termination<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"gopkg.in/natefinch/lumberjack.v2"
	"k8s.io/klog"
)

func main() {
	os.Exit(run())
}

func run() int {
	terminationLog := flag.String("termination-log-file", "", "Write logs after SIGTERM to this file (in addition to stderr)")
	terminationLock := flag.String("termination-touch-file", "", "Touch this file on SIGTERM and delete on termination")

	klog.InitFlags(nil)
	flag.Set("v", "9")

	// never log to stderr, only through our termination log writer (which sends it also to stderr)
	flag.Set("logtostderr", "false")
	flag.Set("stderrthreshold", "99")

	flag.Parse()
	args := flag.CommandLine.Args()

	if len(args) == 0 {
		fmt.Println("Missing command line")
		return 1
	}

	// use special tee-like writer when termination log is set
	termCh := make(chan struct{})
	var stderr io.Writer = os.Stderr
	if len(*terminationLog) > 0 {
		stderr = &terminationFileWriter{
			Writer:             os.Stderr,
			fn:                 *terminationLog,
			startFileLoggingCh: termCh,
		}

		// do the klog file writer dance: klog writes to all outputs of lower
		// severity. No idea why. So we discard for anything other than info.
		// Otherwise, we would see errors multiple times.
		klog.SetOutput(ioutil.Discard)
		klog.SetOutputBySeverity("INFO", stderr)
	}

	// touch file early. If the file is not removed on termination, we are not
	// terminating cleanly via SIGTERM.
	if len(*terminationLock) > 0 {
		klog.Infof("Touching termination lock file %q", *terminationLock)
		if err := touch(*terminationLock); err != nil {
			klog.Infof("error touching %s: %v", *terminationLock, err)
			// keep going
		}
		defer func() {
			klog.Infof("Deleting termination lock file %q", *terminationLock)
			os.Remove(*terminationLock)
		}()
	}

	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = stderr

	// forward SIGTERM and SIGINT to child
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		for s := range sigCh {
			select {
			case <-termCh:
			default:
				close(termCh)
			}

			klog.Infof("Received signal %s. Forwarding to sub-process %q.", s, args[0])

			cmd.Process.Signal(s)
		}
	}()

	klog.Infof("Launching sub-process %q", cmd)
	rc := 0
	if err := cmd.Run(); err != nil {
		if exitError, ok := err.(*exec.ExitError); ok {
			rc = exitError.ExitCode()
		} else {
			klog.Infof("Failed to launch %s: %v", args[0], err)
			return 255
		}
	}

	// remove signal handling
	signal.Stop(sigCh)
	close(sigCh)
	wg.Wait()

	klog.Infof("Termination finished with exit code %d", rc)
	return rc
}

// terminationFileWriter forwards everything to the embedded writer. When
// startFileLoggingCh is closed, everything is appended to the given file name
// in addition.
type terminationFileWriter struct {
	io.Writer
	fn                 string
	startFileLoggingCh <-chan struct{}

	logger io.Writer
}

func (w *terminationFileWriter) Write(bs []byte) (int, error) {
	// temporary hack to avoid logging sensitive tokens.
	// TODO: drop when we moved to a non-sensitive storage format
	if strings.Contains(string(bs), "URI=\"/apis/oauth.openshift.io/v1/oauthaccesstokens/") || strings.Contains(string(bs), "URI=\"/apis/oauth.openshift.io/v1/oauthauthorizetokens/") {
		return len(bs), nil
	}

	select {
	case <-w.startFileLoggingCh:
		if w.logger == nil {
			l := &lumberjack.Logger{
				Filename:   w.fn,
				MaxSize:    100,
				MaxBackups: 3,
				MaxAge:     28,
				Compress:   false,
			}
			w.logger = l
			fmt.Fprintf(os.Stderr, "Copying termination logs to %q\n", w.fn)
		}
		if n, err := w.logger.Write(bs); err != nil {
			return n, err
		} else if n != len(bs) {
			return n, io.ErrShortWrite
		}
	default:
	}

	return w.Writer.Write(bs)
}

func touch(fn string) error {
	_, err := os.Stat(fn)
	if os.IsNotExist(err) {
		file, err := os.Create(fn)
		if err != nil {
			return err
		}
		defer file.Close()
		return nil
	}

	currentTime := time.Now().Local()
	return os.Chtimes(fn, currentTime, currentTime)
}
<|endoftext|>


<commit_before>package scheduler

import (
	"runtime"
	"sort"
	"sync"
	"sync/atomic"
	// "time"
	// "fmt"

	"github.com/henrylee2cn/pholcus/app/aid/history"
	"github.com/henrylee2cn/pholcus/app/aid/proxy"
	"github.com/henrylee2cn/pholcus/app/downloader/context"
	"github.com/henrylee2cn/pholcus/common/util"
	"github.com/henrylee2cn/pholcus/logs"
	"github.com/henrylee2cn/pholcus/runtime/cache"
	"github.com/henrylee2cn/pholcus/runtime/status"
)

type scheduler struct {
	// Spider实例的请求矩阵列表
	matrices map[int]*Matrix
	// 总并发量计数
	count chan bool
	// 运行状态
	status int
	// 全局代理IP
	proxy *proxy.Proxy
	// 标记是否使用代理IP
	useProxy bool
	// 全局历史记录
	history history.Historier
	// 全局读写锁
	sync.RWMutex
}

// 定义全局调度
var sdl = &scheduler{
	history: history.New(),
	status:  status.RUN,
	count:   make(chan bool, cache.Task.ThreadNum),
	proxy:   proxy.New(),
}

func Init() {
	for sdl.proxy == nil {
		runtime.Gosched()
	}
	sdl.matrices = make(map[int]*Matrix)
	sdl.count = make(chan bool, cache.Task.ThreadNum)
	if cache.Task.Mode == status.OFFLINE {
		sdl.history.ReadSuccess(cache.Task.OutType, cache.Task.SuccessInherit)
		sdl.history.ReadFailure(cache.Task.OutType, cache.Task.FailureInherit)
	}
	if cache.Task.ProxyMinute > 0 {
		if sdl.proxy.Count() > 0 {
			sdl.useProxy = true
			sdl.proxy.UpdateTicker(cache.Task.ProxyMinute)
			logs.Log.Informational(" *     使用代理IP，代理IP更换频率为 %v 分钟\n", cache.Task.ProxyMinute)
		} else {
			sdl.useProxy = false
			logs.Log.Informational(" *     在线代理IP列表为空，无法使用代理IP\n")
		}
	} else {
		sdl.useProxy = false
		logs.Log.Informational(" *     不使用代理IP\n")
	}
	sdl.status = status.RUN
}

// 暂停\恢复所有爬行任务
func PauseRecover() {
	sdl.Lock()
	defer sdl.Unlock()
	switch sdl.status {
	case status.PAUSE:
		sdl.status = status.RUN
	case status.RUN:
		sdl.status = status.PAUSE
	}
}

// 终止任务
func Stop() {
	sdl.Lock()
	defer func() {
		recover()
		sdl.Unlock()
	}()

	sdl.status = status.STOP
	for _, v := range sdl.matrices {
		for _, vv := range v.reqs {
			for _, req := range vv {
				// 删除队列中未执行请求的去重记录
				DeleteSuccess(req)
			}
		}
		v.reqs = make(map[int][]*context.Request)
		v.priorities = []int{}
	}

	// 清空
	close(sdl.count)
	sdl.matrices = make(map[int]*Matrix)
}

func UpsertSuccess(record history.Record) bool {
	return sdl.history.UpsertSuccess(record)
}

func DeleteSuccess(record history.Record) {
	sdl.history.DeleteSuccess(record)
}

func UpsertFailure(req *context.Request) bool {
	return sdl.history.UpsertFailure(req)
}

func DeleteFailure(req *context.Request) {
	sdl.history.DeleteFailure(req)
}

// 获取指定蜘蛛在上一次运行时失败的请求
func PullFailure(spiderName string) []*context.Request {
	return sdl.history.PullFailure(spiderName)
}

func TryFlushHistory() {
	if cache.Task.SuccessInherit {
		sdl.history.FlushSuccess(cache.Task.OutType)
	}
	if cache.Task.FailureInherit {
		sdl.history.FlushFailure(cache.Task.OutType)
	}
}

// 一个Spider实例的请求矩阵
type Matrix struct {
	// [优先级]队列，优先级默认为0
	reqs map[int][]*context.Request
	// 优先级顺序，从低到高
	priorities []int
	// 资源使用情况计数
	resCount int32
	// 最大采集页数，以负数形式表示
	maxPage int64
	// 历史及本次失败请求
	failures map[string]*context.Request
	sync.Mutex
}

// 注册资源队列
func NewMatrix(spiderId int, maxPage int64) *Matrix {
	sdl.Lock()
	defer sdl.Unlock()
	matrix := &Matrix{
		reqs:       make(map[int][]*context.Request),
		priorities: []int{},
		maxPage:    maxPage,
		failures:   make(map[string]*context.Request),
	}
	sdl.matrices[spiderId] = matrix
	return matrix
}

// 添加请求到队列
func (self *Matrix) Push(req *context.Request) {
	sdl.RLock()
	defer sdl.RUnlock()

	if sdl.status == status.STOP ||
		self.maxPage >= 0 ||
		// 当req不可重复下载时，已存在成功记录则返回
		!req.IsReloadable() && !UpsertSuccess(req) {
		return
	}

	// 大致限制加入队列的请求量，并发情况下应该会比maxPage多
	atomic.AddInt64(&self.maxPage, 1)

	priority := req.GetPriority()

	// 初始化该蜘蛛下该优先级队列
	if _, found := self.reqs[priority]; !found {
		self.priorities = append(self.priorities, priority)
		sort.Ints(self.priorities) // 从小到大排序
		self.reqs[priority] = []*context.Request{}
	}

	// 添加请求到队列
	self.reqs[priority] = append(self.reqs[priority], req)
}

func (self *Matrix) Pull() (req *context.Request) {
	sdl.RLock()
	defer sdl.RUnlock()
	if sdl.status != status.RUN {
		return
	}

	// 按优先级从高到低取出请求
	for i := len(self.reqs) - 1; i >= 0; i-- {
		idx := self.priorities[i]
		if len(self.reqs[idx]) > 0 {
			req = self.reqs[idx][0]
			self.reqs[idx] = self.reqs[idx][1:]
			if sdl.useProxy {
				req.SetProxy(sdl.proxy.GetOne(req.GetUrl()))
			} else {
				req.SetProxy("")
			}
			return
		}
	}
	return
}

func (self *Matrix) Len() int {
	var l int
	for _, reqs := range self.reqs {
		l += len(reqs)
	}
	return l
}

func (self *Matrix) Use() {
	defer func() {
		recover()
	}()
	sdl.count <- true
	atomic.AddInt32(&self.resCount, 1)
}

func (self *Matrix) Free() {
	<-sdl.count
	atomic.AddInt32(&self.resCount, -1)
}

func (self *Matrix) CanStop() bool {
	sdl.RLock()
	if sdl.status == status.STOP {
		sdl.RUnlock()
		return true
	}
	sdl.RUnlock()

	if self.maxPage >= 0 {
		return true
	}

	if self.resCount != 0 {
		return false
	}

	for i, count := 0, len(self.reqs); i < count; i++ {
		if len(self.reqs[i]) > 0 {
			return false
		}
	}

	if len(self.failures) > 0 {
		// 重新下载历史记录中失败的请求
		var goon bool
		for unique, req := range self.failures {
			if req == nil {
				continue
			}
			self.failures[unique] = nil
			goon = true
			logs.Log.Informational(" *     - 失败请求: [%v]\n", req.GetUrl())
			self.Push(req)
		}
		if goon {
			return false
		}
	}
	return true
}

// 等待处理中的请求完成
func (self *Matrix) Flush() {
	for self.resCount != 0 {
		runtime.Gosched()
		// time.Sleep(5e8)
	}
}

func (self *Matrix) SetFailures(reqs []*context.Request) {
	for _, req := range reqs {
		self.failures[makeUnique(req)] = req
		logs.Log.Informational(" *     + 失败请求: [%v]\n", req.GetUrl())
	}
}

func (self *Matrix) SetFailure(req *context.Request) bool {
	self.Lock()
	defer self.Unlock()
	unique := makeUnique(req)
	if _, ok := self.failures[unique]; !ok {
		// 首次失败时，在任务队列末尾重新执行一次
		self.failures[unique] = req
		logs.Log.Informational(" *     + 失败请求: [%v]\n", req.GetUrl())
		return true
	}
	// 失败两次后，加入历史失败记录
	UpsertFailure(req)
	return false
}

func makeUnique(req *context.Request) string {
	return util.MakeUnique(req.GetUrl() + req.GetMethod())
}
<commit_msg>优化sdl.Push(),降低内存占用<commit_after>package scheduler

import (
	"runtime"
	"sort"
	"sync"
	"sync/atomic"
	// "time"
	// "fmt"

	"github.com/henrylee2cn/pholcus/app/aid/history"
	"github.com/henrylee2cn/pholcus/app/aid/proxy"
	"github.com/henrylee2cn/pholcus/app/downloader/context"
	"github.com/henrylee2cn/pholcus/common/util"
	"github.com/henrylee2cn/pholcus/logs"
	"github.com/henrylee2cn/pholcus/runtime/cache"
	"github.com/henrylee2cn/pholcus/runtime/status"
)

type scheduler struct {
	// Spider实例的请求矩阵列表
	matrices map[int]*Matrix
	// 总并发量计数
	count chan bool
	// 运行状态
	status int
	// 全局代理IP
	proxy *proxy.Proxy
	// 标记是否使用代理IP
	useProxy bool
	// 全局历史记录
	history history.Historier
	// 全局读写锁
	sync.RWMutex
}

// 定义全局调度
var sdl = &scheduler{
	history: history.New(),
	status:  status.RUN,
	count:   make(chan bool, cache.Task.ThreadNum),
	proxy:   proxy.New(),
}

func Init() {
	for sdl.proxy == nil {
		runtime.Gosched()
	}
	sdl.matrices = make(map[int]*Matrix)
	sdl.count = make(chan bool, cache.Task.ThreadNum)
	if cache.Task.Mode == status.OFFLINE {
		sdl.history.ReadSuccess(cache.Task.OutType, cache.Task.SuccessInherit)
		sdl.history.ReadFailure(cache.Task.OutType, cache.Task.FailureInherit)
	}
	if cache.Task.ProxyMinute > 0 {
		if sdl.proxy.Count() > 0 {
			sdl.useProxy = true
			sdl.proxy.UpdateTicker(cache.Task.ProxyMinute)
			logs.Log.Informational(" *     使用代理IP，代理IP更换频率为 %v 分钟\n", cache.Task.ProxyMinute)
		} else {
			sdl.useProxy = false
			logs.Log.Informational(" *     在线代理IP列表为空，无法使用代理IP\n")
		}
	} else {
		sdl.useProxy = false
		logs.Log.Informational(" *     不使用代理IP\n")
	}
	sdl.status = status.RUN
}

// 暂停\恢复所有爬行任务
func PauseRecover() {
	sdl.Lock()
	defer sdl.Unlock()
	switch sdl.status {
	case status.PAUSE:
		sdl.status = status.RUN
	case status.RUN:
		sdl.status = status.PAUSE
	}
}

// 终止任务
func Stop() {
	sdl.Lock()
	defer func() {
		recover()
		sdl.Unlock()
	}()

	sdl.status = status.STOP
	for _, v := range sdl.matrices {
		for _, vv := range v.reqs {
			for _, req := range vv {
				// 删除队列中未执行请求的去重记录
				DeleteSuccess(req)
			}
		}
		v.reqs = make(map[int][]*context.Request)
		v.priorities = []int{}
	}

	// 清空
	close(sdl.count)
	sdl.matrices = make(map[int]*Matrix)
}

func UpsertSuccess(record history.Record) bool {
	return sdl.history.UpsertSuccess(record)
}

func DeleteSuccess(record history.Record) {
	sdl.history.DeleteSuccess(record)
}

func UpsertFailure(req *context.Request) bool {
	return sdl.history.UpsertFailure(req)
}

func DeleteFailure(req *context.Request) {
	sdl.history.DeleteFailure(req)
}

// 获取指定蜘蛛在上一次运行时失败的请求
func PullFailure(spiderName string) []*context.Request {
	return sdl.history.PullFailure(spiderName)
}

func TryFlushHistory() {
	if cache.Task.SuccessInherit {
		sdl.history.FlushSuccess(cache.Task.OutType)
	}
	if cache.Task.FailureInherit {
		sdl.history.FlushFailure(cache.Task.OutType)
	}
}

// 一个Spider实例的请求矩阵
type Matrix struct {
	// [优先级]队列，优先级默认为0
	reqs map[int][]*context.Request
	// 优先级顺序，从低到高
	priorities []int
	// 资源使用情况计数
	resCount int32
	// 最大采集页数，以负数形式表示
	maxPage int64
	// 历史及本次失败请求
	failures map[string]*context.Request
	sync.Mutex
}

// 注册资源队列
func NewMatrix(spiderId int, maxPage int64) *Matrix {
	sdl.Lock()
	defer sdl.Unlock()
	matrix := &Matrix{
		reqs:       make(map[int][]*context.Request),
		priorities: []int{},
		maxPage:    maxPage,
		failures:   make(map[string]*context.Request),
	}
	sdl.matrices[spiderId] = matrix
	return matrix
}

// 添加请求到队列
func (self *Matrix) Push(req *context.Request) {
	// 暂停添加操作，避免大量请求堆积
	for sdl.status == status.PAUSE {
		runtime.Gosched()
	}

	sdl.RLock()
	defer sdl.RUnlock()

	// 终止添加操作
	if sdl.status == status.STOP ||
		self.maxPage >= 0 ||
		// 当req不可重复下载时，已存在成功记录则返回
		!req.IsReloadable() && !UpsertSuccess(req) {
		return
	}

	// 根据资源使用情况，一定程度避免大量请求堆积
	avg := int32(cap(sdl.count) / len(sdl.matrices))
	if avg == 0 {
		avg = 1
	}
	for self.resCount > avg {
		runtime.Gosched()
	}

	// 大致限制加入队列的请求量，并发情况下应该会比maxPage多
	atomic.AddInt64(&self.maxPage, 1)

	priority := req.GetPriority()

	// 初始化该蜘蛛下该优先级队列
	if _, found := self.reqs[priority]; !found {
		self.priorities = append(self.priorities, priority)
		sort.Ints(self.priorities) // 从小到大排序
		self.reqs[priority] = []*context.Request{}
	}

	// 添加请求到队列
	self.reqs[priority] = append(self.reqs[priority], req)
}

// 从队列取出请求，不存在时返回nil
func (self *Matrix) Pull() (req *context.Request) {
	sdl.RLock()
	defer sdl.RUnlock()
	if sdl.status != status.RUN {
		return
	}

	// 按优先级从高到低取出请求
	for i := len(self.reqs) - 1; i >= 0; i-- {
		idx := self.priorities[i]
		if len(self.reqs[idx]) > 0 {
			req = self.reqs[idx][0]
			self.reqs[idx] = self.reqs[idx][1:]
			if sdl.useProxy {
				req.SetProxy(sdl.proxy.GetOne(req.GetUrl()))
			} else {
				req.SetProxy("")
			}
			return
		}
	}
	return
}

func (self *Matrix) Len() int {
	var l int
	for _, reqs := range self.reqs {
		l += len(reqs)
	}
	return l
}

func (self *Matrix) Use() {
	defer func() {
		recover()
	}()
	sdl.count <- true
	atomic.AddInt32(&self.resCount, 1)
}

func (self *Matrix) Free() {
	<-sdl.count
	atomic.AddInt32(&self.resCount, -1)
}

func (self *Matrix) CanStop() bool {
	sdl.RLock()
	if sdl.status == status.STOP {
		sdl.RUnlock()
		return true
	}
	sdl.RUnlock()

	if self.maxPage >= 0 {
		return true
	}

	if self.resCount != 0 {
		return false
	}

	for i, count := 0, len(self.reqs); i < count; i++ {
		if len(self.reqs[i]) > 0 {
			return false
		}
	}

	if len(self.failures) > 0 {
		// 重新下载历史记录中失败的请求
		var goon bool
		for unique, req := range self.failures {
			if req == nil {
				continue
			}
			self.failures[unique] = nil
			goon = true
			logs.Log.Informational(" *     - 失败请求: [%v]\n", req.GetUrl())
			self.Push(req)
		}
		if goon {
			return false
		}
	}
	return true
}

// 等待处理中的请求完成
func (self *Matrix) Flush() {
	for self.resCount != 0 {
		runtime.Gosched()
		// time.Sleep(5e8)
	}
}

func (self *Matrix) SetFailures(reqs []*context.Request) {
	for _, req := range reqs {
		self.failures[makeUnique(req)] = req
		logs.Log.Informational(" *     + 失败请求: [%v]\n", req.GetUrl())
	}
}

func (self *Matrix) SetFailure(req *context.Request) bool {
	self.Lock()
	defer self.Unlock()
	unique := makeUnique(req)
	if _, ok := self.failures[unique]; !ok {
		// 首次失败时，在任务队列末尾重新执行一次
		self.failures[unique] = req
		logs.Log.Informational(" *     + 失败请求: [%v]\n", req.GetUrl())
		return true
	}
	// 失败两次后，加入历史失败记录
	UpsertFailure(req)
	return false
}

func makeUnique(req *context.Request) string {
	return util.MakeUnique(req.GetUrl() + req.GetMethod())
}
<|endoftext|>


<commit_before>package queue

import (
	"crypto/rand"
	"errors"
	"sync"
	"time"
)

const MaxReserTime = 30

var ERROR_HASH_NOT_FOUND = errors.New("Hash not found")
var ERROR_ALREADY_RESERVED = errors.New("Item already reserved")

// Item struct is the basic queue item
type Item struct {
	ReservedAt time.Time `json:"reserved_at"`
	Tag        string    `json:"tag"`
	Value      []byte    `json:"value"`
}

// Data contains all data in the current queue
type Data struct {
	Tag      string          `json:"tag"`
	ItemList map[string]Item `json:"item"`
	control  queueControl
}

type queueControl struct {
	mutex *sync.Mutex
}

// New queue
func New() (q *Data, err error) {
	q = &Data{
		ItemList: make(map[string]Item),
		control:  queueControl{mutex: &sync.Mutex{}},
	}
	return
}

// Count items in the queue
func (q *Data) Count() int {
	return len(q.ItemList)
}

// Put data in the queue
func (q *Data) Put(b []byte) {
	q.ItemList[randStr()] = Item{Value: b}
}

// Reserve an item to be processed.
// If the item is not removed or the reservation time is not
// renewed, the item will returns to the queue automatically
func (q *Data) Reserve(hash string) (item Item, err error) {
	q.control.mutex.Lock()
	defer q.control.mutex.Unlock()
	v, ok := q.ItemList[hash]
	if !ok {
		err = ERROR_HASH_NOT_FOUND
		return
	}
	now := time.Now()
	diff := now.Sub(v.ReservedAt)
	if diff.Seconds() < MaxReserTime {
		err = ERROR_ALREADY_RESERVED
		return
	}
	v.ReservedAt = time.Now()
	item = v
	return
}

func randStr() (ret string) {
	const charList = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
	const clLen byte = 62
	const size = 36
	var bytes = make([]byte, size)
	if _, err := rand.Read(bytes); err != nil {
		panic(err.Error())
	}
	for k, v := range bytes {
		bytes[k] = charList[v%clLen]
	}
	ret = string(bytes)
	return
}
<commit_msg>remove implemented<commit_after>package queue

import (
	"crypto/rand"
	"errors"
	"sync"
	"time"
)

const MaxReserTime = 30

var ERROR_HASH_NOT_FOUND = errors.New("Hash not found")
var ERROR_ALREADY_RESERVED = errors.New("Item already reserved")
var ERROR_NOT_RESERVED = errors.New("Item not reserved")

// Item struct is the basic queue item
type Item struct {
	ReservedAt time.Time `json:"reserved_at"`
	Tag        string    `json:"tag"`
	Value      []byte    `json:"value"`
}

// Data contains all data in the current queue
type Data struct {
	Tag      string          `json:"tag"`
	ItemList map[string]Item `json:"item"`
	control  queueControl
}

type queueControl struct {
	mutex *sync.Mutex
}

// New queue
func New() (q *Data, err error) {
	q = &Data{
		ItemList: make(map[string]Item),
		control:  queueControl{mutex: &sync.Mutex{}},
	}
	return
}

// Count items in the queue
func (q *Data) Count() int {
	return len(q.ItemList)
}

// Put data in the queue
func (q *Data) Put(b []byte) {
	q.ItemList[randStr()] = Item{Value: b}
}

// Reserve an item to be processed.
// If the item is not removed or the reservation time is not
// renewed, the item will returns to the queue automatically
func (q *Data) Reserve(hash string) (item Item, err error) {
	q.control.mutex.Lock()
	defer q.control.mutex.Unlock()
	v, ok := q.ItemList[hash]
	if !ok {
		err = ERROR_HASH_NOT_FOUND
		return
	}
	now := time.Now()
	diff := now.Sub(v.ReservedAt)
	if diff.Seconds() < MaxReserTime {
		err = ERROR_ALREADY_RESERVED
		return
	}
	v.ReservedAt = time.Now()
	item = v
	return
}

// Remove item from the queue, the item must be reserved.
func (q *Data) Remove(hash string) (err error) {
	q.control.mutex.Lock()
	defer q.control.mutex.Unlock()
	v, ok := q.ItemList[hash]
	if !ok {
		err = ERROR_HASH_NOT_FOUND
		return
	}
	now := time.Now()
	diff := now.Sub(v.ReservedAt)
	if diff.Seconds() >= MaxReserTime {
		err = ERROR_NOT_RESERVED
		return
	}
	delete(q.ItemList, hash)
	return
}

func randStr() (ret string) {
	const charList = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
	const clLen byte = 62
	const size = 36
	var bytes = make([]byte, size)
	if _, err := rand.Read(bytes); err != nil {
		panic(err.Error())
	}
	for k, v := range bytes {
		bytes[k] = charList[v%clLen]
	}
	ret = string(bytes)
	return
}
<|endoftext|>


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

package sdk

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"text/template"

	"github.com/coreos/pkg/capnslog"

	"github.com/coreos/mantle/system/exec"
	"github.com/coreos/mantle/system/user"
	"github.com/coreos/mantle/util"
)

// Must run inside the SDK chroot, easiest to just assemble a script to do it
const (
	safePath   = "PATH=/usr/sbin:/usr/bin:/sbin:/bin"
	sudoPrompt = "sudo password for %p: "
	script     = `#!/bin/bash
set -e

# make sure user/group database files exist
touch /etc/{group,gshadow,passwd,shadow}
chmod 0640 /etc/{gshadow,shadow}

# add group if it doesn't exist already
if ! getent group {{printf "%q" .Groupname}} >/dev/null; then
	echo Adding group {{printf "%q" .Groupname}}
	groupadd -o -g {{.Gid}} {{printf "%q" .Groupname}}
fi

# add user if it doesn't exist already
if ! getent passwd {{printf "%q" .Username}} >/dev/null; then
	echo Adding user {{printf "%q" .Username}}
	useradd -o -g {{.Gid}} -u {{.Uid}} -s /bin/bash -m \
		-c {{printf "%q" .Name}} {{printf "%q" .Username}}
fi

for g in kvm portage sudo; do
	# copy system group from /usr to /etc if needed
	if getent -s usrfiles group "$g" >/dev/null && \
	   ! getent -s files group "$g" >/dev/null; then
		getent -s usrfiles group "$g" >> /etc/group
	fi
	gpasswd -a {{printf "%q" .Username}} "$g"
done

echo Setting up sudoers
cat >/etc/sudoers.d/90_env_keep <<EOF
Defaults env_keep += "\
COREOS_BUILD_ID COREOS_OFFICIAL \
EMAIL GIT_AUTHOR_EMAIL GIT_AUTHOR_NAME \
GIT_COMMITTER_EMAIL GIT_COMMITTER_NAME \
GIT_PROXY_COMMAND GIT_SSH RSYNC_PROXY \
GNUPGHOME GPG_AGENT_INFO SSH_AUTH_SOCK \
BOTO_PATH GOOGLE_APPLICATION_CREDENTIALS \
USE FEATURES PORTAGE_USERNAME \
all_proxy ftp_proxy http_proxy https_proxy no_proxy"
EOF
chmod 0440 /etc/sudoers.d/90_env_keep

echo Setting default enviornment variables
cat >/etc/env.d/90portage_username <<EOF
PORTAGE_USERNAME={{printf "%q" .Username}}
EOF
# needlessly noisy since portage isn't set up yet
env-update &>/dev/null

echo Setting up home directory
HOME=/home/{{printf "%q" .Username}}

# Create ~/trunk symlink
ln -sfT /mnt/host/source "$HOME"/trunk

rm -f "$HOME"/.bash{_logout,_profile,rc}
cat >"$HOME"/.bash_logout <<EOF
# .bash_logout

# This file is sourced when a login shell terminates.
EOF

cat >"$HOME"/.bash_profile <<EOF
# .bash_profile

# This file is sourced by bash for login shells.  The following line
# runs your .bashrc and is recommended by the bash info pages.
[[ -f ~/.bashrc ]] && . ~/.bashrc

# Automatically change to scripts directory.
cd ${CHROOT_CWD:-~/trunk/src/scripts}
EOF

cat >"$HOME"/.bashrc <<EOF
# .bashrc

# This file is sourced by all *interactive* bash shells on startup,
# including some apparently interactive shells such as scp and rcp
# that can't tolerate any output.  So make sure this doesn't display
# anything or bad things will happen !

# Test for an interactive shell.  There is no need to set anything
# past this point for scp and rcp, and it's important to refrain from
# outputting anything in those cases.
if [[ $- != *i* ]] ; then
	# Shell is non-interactive.  Be done now!
	return
fi

# Enable bash completion for build scripts.
source ~/trunk/src/scripts/bash_completion

# Put your fun stuff here.
EOF

chown -R {{.Uid}}:{{.Gid}} "$HOME"

# Checked in src/scripts/common.sh
touch /etc/debian_chroot
`
)

var scriptTemplate = template.Must(template.New("script").Parse(script))

func Setup(name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	u, err := user.Current()
	if err != nil {
		return err
	}

	var sc bytes.Buffer
	if err := scriptTemplate.Execute(&sc, u); err != nil {
		return err
	}

	plog.Info("Configuring SDK chroot")
	sh := exec.Command(
		"sudo", "-p", sudoPrompt,
		"chroot", chroot,
		"/usr/bin/env", "-i",
		"/bin/bash", "--login")
	sh.Stdin = &sc
	sh.Stderr = os.Stderr
	if plog.LevelAt(capnslog.INFO) {
		out, err := sh.StdoutPipe()
		if err != nil {
			return err
		}
		go util.LogFrom(capnslog.INFO, out)
	}
	if plog.LevelAt(capnslog.DEBUG) {
		sh.Args = append(sh.Args, "-x")
	}
	return sh.Run()
}

func extract(tar, dir string) error {
	in, err := os.Open(tar)
	if err != nil {
		return err
	}
	defer in.Close()

	unzipper, err := exec.LookPath("lbzcat")
	if err != nil {
		unzipper = "bzcat"
	}

	unzip := exec.Command(unzipper)
	unzip.Stdin = in
	unzip.Stderr = os.Stderr
	unzipped, err := unzip.StdoutPipe()
	if err != nil {
		return err
	}

	untar := exec.Command("sudo", "-p", sudoPrompt,
		"tar", "--numeric-owner", "-x")
	untar.Dir = dir
	untar.Stdin = unzipped
	untar.Stderr = os.Stderr

	if err := unzip.Start(); err != nil {
		return err
	}

	if err := untar.Start(); err != nil {
		unzip.Kill()
		return err
	}

	if err := untar.Wait(); err != nil {
		unzip.Kill()
		return err
	}

	if err := unzip.Wait(); err != nil {
		return err
	}

	return nil
}

func Unpack(version, name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	if _, err := os.Stat(chroot); !os.IsNotExist(err) {
		if err == nil {
			err = fmt.Errorf("Path already exists: %s", chroot)
		}
		return err
	}

	plog.Noticef("Unpacking SDK into %s", chroot)
	if err := os.MkdirAll(chroot, 0777); err != nil {
		return err
	}

	tar := filepath.Join(RepoCache(), "sdks", TarballName(version))
	plog.Infof("Using %s", tar)
	if err := extract(tar, chroot); err != nil {
		plog.Errorf("Extracting %s to %s failed: %v", tar, chroot, err)
		return err
	}
	plog.Notice("Unpacked")

	return nil
}

func Delete(name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	if _, err := os.Stat(chroot); err != nil {
		if os.IsNotExist(err) {
			plog.Infof("Path does not exist: %s", chroot)
			return nil
		}
		return err
	}

	plog.Noticef("Removing SDK at %s", chroot)
	rm := exec.Command("sudo", "-p", sudoPrompt, "rm", "-rf", chroot)
	rm.Stderr = os.Stderr
	if err := rm.Run(); err != nil {
		return err
	}
	plog.Notice("Removed")

	return nil
}
<commit_msg>cork/create: fix interactive shell check<commit_after>// Copyright 2015 CoreOS, Inc.
// Copyright 2011 The Go Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package sdk

import (
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"text/template"

	"github.com/coreos/pkg/capnslog"

	"github.com/coreos/mantle/system/exec"
	"github.com/coreos/mantle/system/user"
	"github.com/coreos/mantle/util"
)

// Must run inside the SDK chroot, easiest to just assemble a script to do it
const (
	safePath   = "PATH=/usr/sbin:/usr/bin:/sbin:/bin"
	sudoPrompt = "sudo password for %p: "
	script     = `#!/bin/bash
set -e

# make sure user/group database files exist
touch /etc/{group,gshadow,passwd,shadow}
chmod 0640 /etc/{gshadow,shadow}

# add group if it doesn't exist already
if ! getent group {{printf "%q" .Groupname}} >/dev/null; then
	echo Adding group {{printf "%q" .Groupname}}
	groupadd -o -g {{.Gid}} {{printf "%q" .Groupname}}
fi

# add user if it doesn't exist already
if ! getent passwd {{printf "%q" .Username}} >/dev/null; then
	echo Adding user {{printf "%q" .Username}}
	useradd -o -g {{.Gid}} -u {{.Uid}} -s /bin/bash -m \
		-c {{printf "%q" .Name}} {{printf "%q" .Username}}
fi

for g in kvm portage sudo; do
	# copy system group from /usr to /etc if needed
	if getent -s usrfiles group "$g" >/dev/null && \
	   ! getent -s files group "$g" >/dev/null; then
		getent -s usrfiles group "$g" >> /etc/group
	fi
	gpasswd -a {{printf "%q" .Username}} "$g"
done

echo Setting up sudoers
cat >/etc/sudoers.d/90_env_keep <<EOF
Defaults env_keep += "\
COREOS_BUILD_ID COREOS_OFFICIAL \
EMAIL GIT_AUTHOR_EMAIL GIT_AUTHOR_NAME \
GIT_COMMITTER_EMAIL GIT_COMMITTER_NAME \
GIT_PROXY_COMMAND GIT_SSH RSYNC_PROXY \
GNUPGHOME GPG_AGENT_INFO SSH_AUTH_SOCK \
BOTO_PATH GOOGLE_APPLICATION_CREDENTIALS \
USE FEATURES PORTAGE_USERNAME \
all_proxy ftp_proxy http_proxy https_proxy no_proxy"
EOF
chmod 0440 /etc/sudoers.d/90_env_keep

echo Setting default enviornment variables
cat >/etc/env.d/90portage_username <<EOF
PORTAGE_USERNAME={{printf "%q" .Username}}
EOF
# needlessly noisy since portage isn't set up yet
env-update &>/dev/null

echo Setting up home directory
HOME=/home/{{printf "%q" .Username}}

# Create ~/trunk symlink
ln -sfT /mnt/host/source "$HOME"/trunk

rm -f "$HOME"/.bash{_logout,_profile,rc}
cat >"$HOME"/.bash_logout <<EOF
# .bash_logout

# This file is sourced when a login shell terminates.
EOF

cat >"$HOME"/.bash_profile <<EOF
# .bash_profile

# This file is sourced by bash for login shells.  The following line
# runs your .bashrc and is recommended by the bash info pages.
[[ -f ~/.bashrc ]] && . ~/.bashrc

# Automatically change to scripts directory.
cd ${CHROOT_CWD:-~/trunk/src/scripts}
EOF

cat >"$HOME"/.bashrc << 'EOF'
# .bashrc

# This file is sourced by all *interactive* bash shells on startup,
# including some apparently interactive shells such as scp and rcp
# that can't tolerate any output.  So make sure this doesn't display
# anything or bad things will happen !

# Test for an interactive shell.  There is no need to set anything
# past this point for scp and rcp, and it's important to refrain from
# outputting anything in those cases.
if [[ $- != *i* ]] ; then
	# Shell is non-interactive.  Be done now!
	return
fi

# Enable bash completion for build scripts.
source ~/trunk/src/scripts/bash_completion

# Put your fun stuff here.
EOF

chown -R {{.Uid}}:{{.Gid}} "$HOME"

# Checked in src/scripts/common.sh
touch /etc/debian_chroot
`
)

var scriptTemplate = template.Must(template.New("script").Parse(script))

func Setup(name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	u, err := user.Current()
	if err != nil {
		return err
	}

	var sc bytes.Buffer
	if err := scriptTemplate.Execute(&sc, u); err != nil {
		return err
	}

	plog.Info("Configuring SDK chroot")
	sh := exec.Command(
		"sudo", "-p", sudoPrompt,
		"chroot", chroot,
		"/usr/bin/env", "-i",
		"/bin/bash", "--login")
	sh.Stdin = &sc
	sh.Stderr = os.Stderr
	if plog.LevelAt(capnslog.INFO) {
		out, err := sh.StdoutPipe()
		if err != nil {
			return err
		}
		go util.LogFrom(capnslog.INFO, out)
	}
	if plog.LevelAt(capnslog.DEBUG) {
		sh.Args = append(sh.Args, "-x")
	}
	return sh.Run()
}

func extract(tar, dir string) error {
	in, err := os.Open(tar)
	if err != nil {
		return err
	}
	defer in.Close()

	unzipper, err := exec.LookPath("lbzcat")
	if err != nil {
		unzipper = "bzcat"
	}

	unzip := exec.Command(unzipper)
	unzip.Stdin = in
	unzip.Stderr = os.Stderr
	unzipped, err := unzip.StdoutPipe()
	if err != nil {
		return err
	}

	untar := exec.Command("sudo", "-p", sudoPrompt,
		"tar", "--numeric-owner", "-x")
	untar.Dir = dir
	untar.Stdin = unzipped
	untar.Stderr = os.Stderr

	if err := unzip.Start(); err != nil {
		return err
	}

	if err := untar.Start(); err != nil {
		unzip.Kill()
		return err
	}

	if err := untar.Wait(); err != nil {
		unzip.Kill()
		return err
	}

	if err := unzip.Wait(); err != nil {
		return err
	}

	return nil
}

func Unpack(version, name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	if _, err := os.Stat(chroot); !os.IsNotExist(err) {
		if err == nil {
			err = fmt.Errorf("Path already exists: %s", chroot)
		}
		return err
	}

	plog.Noticef("Unpacking SDK into %s", chroot)
	if err := os.MkdirAll(chroot, 0777); err != nil {
		return err
	}

	tar := filepath.Join(RepoCache(), "sdks", TarballName(version))
	plog.Infof("Using %s", tar)
	if err := extract(tar, chroot); err != nil {
		plog.Errorf("Extracting %s to %s failed: %v", tar, chroot, err)
		return err
	}
	plog.Notice("Unpacked")

	return nil
}

func Delete(name string) error {
	chroot := filepath.Join(RepoRoot(), name)
	if _, err := os.Stat(chroot); err != nil {
		if os.IsNotExist(err) {
			plog.Infof("Path does not exist: %s", chroot)
			return nil
		}
		return err
	}

	plog.Noticef("Removing SDK at %s", chroot)
	rm := exec.Command("sudo", "-p", sudoPrompt, "rm", "-rf", chroot)
	rm.Stderr = os.Stderr
	if err := rm.Run(); err != nil {
		return err
	}
	plog.Notice("Removed")

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Package memdb provides in-memory key/value database implementation.
package memdb

import (
	"math/rand"
	"sync"

	"github.com/syndtr/goleveldb/leveldb/comparer"
	"github.com/syndtr/goleveldb/leveldb/iterator"
	"github.com/syndtr/goleveldb/leveldb/util"
)

var (
	ErrNotFound = util.ErrNotFound
)

const tMaxHeight = 12

type dbIter struct {
	util.BasicReleaser
	p          *DB
	slice      *util.Range
	node       int
	forward    bool
	key, value []byte
}

func (i *dbIter) fill(checkStart, checkLimit bool) bool {
	if i.node != 0 {
		n := i.p.nodeData[i.node]
		m := n + i.p.nodeData[i.node+nKey]
		i.key = i.p.kvData[n:m]
		if i.slice != nil {
			switch {
			case checkLimit && i.slice.Limit != nil && i.p.cmp.Compare(i.key, i.slice.Limit) >= 0:
				fallthrough
			case checkStart && i.slice.Start != nil && i.p.cmp.Compare(i.key, i.slice.Start) < 0:
				i.node = 0
				goto bail
			}
		}
		i.value = i.p.kvData[m : m+i.p.nodeData[i.node+nVal]]
		return true
	}
bail:
	i.key = nil
	i.value = nil
	return false
}

func (i *dbIter) Valid() bool {
	return i.node != 0
}

func (i *dbIter) First() bool {
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Start != nil {
		i.node, _ = i.p.findGE(i.slice.Start, false)
	} else {
		i.node = i.p.nodeData[nNext]
	}
	return i.fill(false, true)
}

func (i *dbIter) Last() bool {
	i.forward = false
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Limit != nil {
		i.node = i.p.findLT(i.slice.Limit)
	} else {
		i.node = i.p.findLast()
	}
	return i.fill(true, false)
}

func (i *dbIter) Seek(key []byte) bool {
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Start != nil && i.p.cmp.Compare(key, i.slice.Start) < 0 {
		key = i.slice.Start
	}
	i.node, _ = i.p.findGE(key, false)
	return i.fill(false, true)
}

func (i *dbIter) Next() bool {
	if i.node == 0 {
		if !i.forward {
			return i.First()
		}
		return false
	}
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	i.node = i.p.nodeData[i.node+nNext]
	return i.fill(false, true)
}

func (i *dbIter) Prev() bool {
	if i.node == 0 {
		if i.forward {
			return i.Last()
		}
		return false
	}
	i.forward = false
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	i.node = i.p.findLT(i.key)
	return i.fill(true, false)
}

func (i *dbIter) Key() []byte {
	return i.key
}

func (i *dbIter) Value() []byte {
	return i.value
}

func (i *dbIter) Error() error { return nil }

const (
	nKV = iota
	nKey
	nVal
	nHeight
	nNext
)

// DB is an in-memory key/value database.
type DB struct {
	cmp comparer.BasicComparer
	rnd *rand.Rand

	mu     sync.RWMutex
	kvData []byte
	// Node data:
	// [0]         : KV offset
	// [1]         : Key length
	// [2]         : Value length
	// [3]         : Height
	// [3..height] : Next nodes
	nodeData  []int
	prevNode  [tMaxHeight]int
	maxHeight int
	n         int
	kvSize    int
}

func (p *DB) randHeight() (h int) {
	const branching = 4
	h = 1
	for h < tMaxHeight && p.rnd.Int()%branching == 0 {
		h++
	}
	return
}

func (p *DB) findGE(key []byte, prev bool) (int, bool) {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		cmp := 1
		if next != 0 {
			o := p.nodeData[next]
			cmp = p.cmp.Compare(p.kvData[o:o+p.nodeData[next+nKey]], key)
		}
		if cmp < 0 {
			// Keep searching in this list
			node = next
		} else {
			if prev {
				p.prevNode[h] = node
			} else if cmp == 0 {
				return next, true
			}
			if h == 0 {
				return next, cmp == 0
			}
			h--
		}
	}
}

func (p *DB) findLT(key []byte) int {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		o := p.nodeData[next]
		if next == 0 || p.cmp.Compare(p.kvData[o:o+p.nodeData[next+nKey]], key) >= 0 {
			if h == 0 {
				break
			}
			h--
		} else {
			node = next
		}
	}
	return node
}

func (p *DB) findLast() int {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		if next == 0 {
			if h == 0 {
				break
			}
			h--
		} else {
			node = next
		}
	}
	return node
}

// Put sets the value for the given key. It overwrites any previous value
// for that key; a DB is not a multi-map.
//
// It is safe to modify the contents of the arguments after Put returns.
func (p *DB) Put(key []byte, value []byte) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if node, exact := p.findGE(key, true); exact {
		kvOffset := len(p.kvData)
		p.kvData = append(p.kvData, key...)
		p.kvData = append(p.kvData, value...)
		p.nodeData[node] = kvOffset
		m := p.nodeData[node+nVal]
		p.nodeData[node+nVal] = len(value)
		p.kvSize += len(value) - m
		return nil
	}

	h := p.randHeight()
	if h > p.maxHeight {
		for i := p.maxHeight; i < h; i++ {
			p.prevNode[i] = 0
		}
		p.maxHeight = h
	}

	kvOffset := len(p.kvData)
	p.kvData = append(p.kvData, key...)
	p.kvData = append(p.kvData, value...)
	// Node
	node := len(p.nodeData)
	p.nodeData = append(p.nodeData, kvOffset, len(key), len(value), h)
	for i, n := range p.prevNode[:h] {
		m := n + 4 + i
		p.nodeData = append(p.nodeData, p.nodeData[m])
		p.nodeData[m] = node
	}

	p.kvSize += len(key) + len(value)
	p.n++
	return nil
}

// Delete deletes the value for the given key. It returns ErrNotFound if
// the DB does not contain the key.
//
// It is safe to modify the contents of the arguments after Delete returns.
func (p *DB) Delete(key []byte) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	node, exact := p.findGE(key, true)
	if !exact {
		return ErrNotFound
	}

	h := p.nodeData[node+nHeight]
	for i, n := range p.prevNode[:h] {
		m := n + 4 + i
		p.nodeData[m] = p.nodeData[p.nodeData[m]+nNext+i]
	}

	p.kvSize -= p.nodeData[node+nKey] + p.nodeData[node+nVal]
	p.n--
	return nil
}

// Contains returns true if the given key are in the DB.
//
// It is safe to modify the contents of the arguments after Contains returns.
func (p *DB) Contains(key []byte) bool {
	p.mu.RLock()
	_, exact := p.findGE(key, false)
	p.mu.RUnlock()
	return exact
}

// Get gets the value for the given key. It returns error.ErrNotFound if the
// DB does not contain the key.
//
// The caller should not modify the contents of the returned slice, but
// it is safe to modify the contents of the argument after Get returns.
func (p *DB) Get(key []byte) (value []byte, err error) {
	p.mu.RLock()
	if node, exact := p.findGE(key, false); exact {
		o := p.nodeData[node] + p.nodeData[node+nKey]
		value = p.kvData[o : o+p.nodeData[node+nVal]]
	} else {
		err = ErrNotFound
	}
	p.mu.RUnlock()
	return
}

// Find finds key/value pair whose key is greater than or equal to the
// given key. It returns ErrNotFound if the table doesn't contain
// such pair.
//
// The caller should not modify the contents of the returned slice, but
// it is safe to modify the contents of the argument after Find returns.
func (p *DB) Find(key []byte) (rkey, value []byte, err error) {
	p.mu.RLock()
	if node, _ := p.findGE(key, false); node != 0 {
		n := p.nodeData[node]
		m := n + p.nodeData[node+nKey]
		rkey = p.kvData[n:m]
		value = p.kvData[m : m+p.nodeData[node+nVal]]
	} else {
		err = ErrNotFound
	}
	p.mu.RUnlock()
	return
}

// NewIterator returns an iterator of the DB.
// The returned iterator is not goroutine-safe, but it is safe to use
// multiple iterators concurrently, with each in a dedicated goroutine.
// It is also safe to use an iterator concurrently with modifying its
// underlying DB. However, the resultant key/value pairs are not guaranteed
// to be a consistent snapshot of the DB at a particular point in time.
//
// Slice allows slicing the iterator to only contains keys in the given
// range. A nil Range.Start is treated as a key before all keys in the
// DB. And a nil Range.Limit is treated as a key after all keys in
// the DB.
//
// The iterator must be released after use, by calling Release method.
//
// Also read Iterator documentation of the leveldb/iterator package.
func (p *DB) NewIterator(slice *util.Range) iterator.Iterator {
	return &dbIter{p: p, slice: slice}
}

// Capacity returns keys/values buffer capacity.
func (p *DB) Capacity() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return cap(p.kvData)
}

// Size returns sum of keys and values length. Note that deleted
// key/value will not be accouted for, but it will still consume
// the buffer, since the buffer is append only.
func (p *DB) Size() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.kvSize
}

// Free returns keys/values free buffer before need to grow.
func (p *DB) Free() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return cap(p.kvData) - len(p.kvData)
}

// Len returns the number of entries in the DB.
func (p *DB) Len() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.n
}

// Reset resets the DB to initial empty state. Allows reuse the buffer.
func (p *DB) Reset() {
	p.rnd = rand.New(rand.NewSource(0xdeadbeef))
	p.maxHeight = 1
	p.n = 0
	p.kvSize = 0
	p.kvData = p.kvData[:0]
	p.nodeData = p.nodeData[:4+tMaxHeight]
	p.nodeData[nKV] = 0
	p.nodeData[nKey] = 0
	p.nodeData[nVal] = 0
	p.nodeData[nHeight] = tMaxHeight
	for n := 0; n < tMaxHeight; n++ {
		p.nodeData[4+n] = 0
		p.prevNode[n] = 0
	}
}

// New creates a new initalized in-memory key/value DB. The capacity
// is the initial key/value buffer capacity. The capacity is advisory,
// not enforced.
func New(cmp comparer.BasicComparer, capacity int) *DB {
	p := &DB{
		cmp:       cmp,
		rnd:       rand.New(rand.NewSource(0xdeadbeef)),
		maxHeight: 1,
		kvData:    make([]byte, 0, capacity),
		nodeData:  make([]int, 4+tMaxHeight),
	}
	p.nodeData[nHeight] = tMaxHeight
	return p
}
<commit_msg>memdb: Set variables to nil when iterator get released<commit_after>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Package memdb provides in-memory key/value database implementation.
package memdb

import (
	"math/rand"
	"sync"

	"github.com/syndtr/goleveldb/leveldb/comparer"
	"github.com/syndtr/goleveldb/leveldb/iterator"
	"github.com/syndtr/goleveldb/leveldb/util"
)

var (
	ErrNotFound = util.ErrNotFound
)

const tMaxHeight = 12

type dbIter struct {
	util.BasicReleaser
	p          *DB
	slice      *util.Range
	node       int
	forward    bool
	key, value []byte
}

func (i *dbIter) fill(checkStart, checkLimit bool) bool {
	if i.node != 0 {
		n := i.p.nodeData[i.node]
		m := n + i.p.nodeData[i.node+nKey]
		i.key = i.p.kvData[n:m]
		if i.slice != nil {
			switch {
			case checkLimit && i.slice.Limit != nil && i.p.cmp.Compare(i.key, i.slice.Limit) >= 0:
				fallthrough
			case checkStart && i.slice.Start != nil && i.p.cmp.Compare(i.key, i.slice.Start) < 0:
				i.node = 0
				goto bail
			}
		}
		i.value = i.p.kvData[m : m+i.p.nodeData[i.node+nVal]]
		return true
	}
bail:
	i.key = nil
	i.value = nil
	return false
}

func (i *dbIter) Valid() bool {
	return i.node != 0
}

func (i *dbIter) First() bool {
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Start != nil {
		i.node, _ = i.p.findGE(i.slice.Start, false)
	} else {
		i.node = i.p.nodeData[nNext]
	}
	return i.fill(false, true)
}

func (i *dbIter) Last() bool {
	if i.p == nil {
		return false
	}
	i.forward = false
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Limit != nil {
		i.node = i.p.findLT(i.slice.Limit)
	} else {
		i.node = i.p.findLast()
	}
	return i.fill(true, false)
}

func (i *dbIter) Seek(key []byte) bool {
	if i.p == nil {
		return false
	}
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	if i.slice != nil && i.slice.Start != nil && i.p.cmp.Compare(key, i.slice.Start) < 0 {
		key = i.slice.Start
	}
	i.node, _ = i.p.findGE(key, false)
	return i.fill(false, true)
}

func (i *dbIter) Next() bool {
	if i.p == nil {
		return false
	}
	if i.node == 0 {
		if !i.forward {
			return i.First()
		}
		return false
	}
	i.forward = true
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	i.node = i.p.nodeData[i.node+nNext]
	return i.fill(false, true)
}

func (i *dbIter) Prev() bool {
	if i.p == nil {
		return false
	}
	if i.node == 0 {
		if i.forward {
			return i.Last()
		}
		return false
	}
	i.forward = false
	i.p.mu.RLock()
	defer i.p.mu.RUnlock()
	i.node = i.p.findLT(i.key)
	return i.fill(true, false)
}

func (i *dbIter) Key() []byte {
	return i.key
}

func (i *dbIter) Value() []byte {
	return i.value
}

func (i *dbIter) Error() error { return nil }

func (i *dbIter) Release() {
	if i.p != nil {
		i.p = nil
		i.node = 0
		i.key = nil
		i.value = nil
		i.BasicReleaser.Release()
	}
}

const (
	nKV = iota
	nKey
	nVal
	nHeight
	nNext
)

// DB is an in-memory key/value database.
type DB struct {
	cmp comparer.BasicComparer
	rnd *rand.Rand

	mu     sync.RWMutex
	kvData []byte
	// Node data:
	// [0]         : KV offset
	// [1]         : Key length
	// [2]         : Value length
	// [3]         : Height
	// [3..height] : Next nodes
	nodeData  []int
	prevNode  [tMaxHeight]int
	maxHeight int
	n         int
	kvSize    int
}

func (p *DB) randHeight() (h int) {
	const branching = 4
	h = 1
	for h < tMaxHeight && p.rnd.Int()%branching == 0 {
		h++
	}
	return
}

func (p *DB) findGE(key []byte, prev bool) (int, bool) {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		cmp := 1
		if next != 0 {
			o := p.nodeData[next]
			cmp = p.cmp.Compare(p.kvData[o:o+p.nodeData[next+nKey]], key)
		}
		if cmp < 0 {
			// Keep searching in this list
			node = next
		} else {
			if prev {
				p.prevNode[h] = node
			} else if cmp == 0 {
				return next, true
			}
			if h == 0 {
				return next, cmp == 0
			}
			h--
		}
	}
}

func (p *DB) findLT(key []byte) int {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		o := p.nodeData[next]
		if next == 0 || p.cmp.Compare(p.kvData[o:o+p.nodeData[next+nKey]], key) >= 0 {
			if h == 0 {
				break
			}
			h--
		} else {
			node = next
		}
	}
	return node
}

func (p *DB) findLast() int {
	node := 0
	h := p.maxHeight - 1
	for {
		next := p.nodeData[node+nNext+h]
		if next == 0 {
			if h == 0 {
				break
			}
			h--
		} else {
			node = next
		}
	}
	return node
}

// Put sets the value for the given key. It overwrites any previous value
// for that key; a DB is not a multi-map.
//
// It is safe to modify the contents of the arguments after Put returns.
func (p *DB) Put(key []byte, value []byte) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	if node, exact := p.findGE(key, true); exact {
		kvOffset := len(p.kvData)
		p.kvData = append(p.kvData, key...)
		p.kvData = append(p.kvData, value...)
		p.nodeData[node] = kvOffset
		m := p.nodeData[node+nVal]
		p.nodeData[node+nVal] = len(value)
		p.kvSize += len(value) - m
		return nil
	}

	h := p.randHeight()
	if h > p.maxHeight {
		for i := p.maxHeight; i < h; i++ {
			p.prevNode[i] = 0
		}
		p.maxHeight = h
	}

	kvOffset := len(p.kvData)
	p.kvData = append(p.kvData, key...)
	p.kvData = append(p.kvData, value...)
	// Node
	node := len(p.nodeData)
	p.nodeData = append(p.nodeData, kvOffset, len(key), len(value), h)
	for i, n := range p.prevNode[:h] {
		m := n + 4 + i
		p.nodeData = append(p.nodeData, p.nodeData[m])
		p.nodeData[m] = node
	}

	p.kvSize += len(key) + len(value)
	p.n++
	return nil
}

// Delete deletes the value for the given key. It returns ErrNotFound if
// the DB does not contain the key.
//
// It is safe to modify the contents of the arguments after Delete returns.
func (p *DB) Delete(key []byte) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	node, exact := p.findGE(key, true)
	if !exact {
		return ErrNotFound
	}

	h := p.nodeData[node+nHeight]
	for i, n := range p.prevNode[:h] {
		m := n + 4 + i
		p.nodeData[m] = p.nodeData[p.nodeData[m]+nNext+i]
	}

	p.kvSize -= p.nodeData[node+nKey] + p.nodeData[node+nVal]
	p.n--
	return nil
}

// Contains returns true if the given key are in the DB.
//
// It is safe to modify the contents of the arguments after Contains returns.
func (p *DB) Contains(key []byte) bool {
	p.mu.RLock()
	_, exact := p.findGE(key, false)
	p.mu.RUnlock()
	return exact
}

// Get gets the value for the given key. It returns error.ErrNotFound if the
// DB does not contain the key.
//
// The caller should not modify the contents of the returned slice, but
// it is safe to modify the contents of the argument after Get returns.
func (p *DB) Get(key []byte) (value []byte, err error) {
	p.mu.RLock()
	if node, exact := p.findGE(key, false); exact {
		o := p.nodeData[node] + p.nodeData[node+nKey]
		value = p.kvData[o : o+p.nodeData[node+nVal]]
	} else {
		err = ErrNotFound
	}
	p.mu.RUnlock()
	return
}

// Find finds key/value pair whose key is greater than or equal to the
// given key. It returns ErrNotFound if the table doesn't contain
// such pair.
//
// The caller should not modify the contents of the returned slice, but
// it is safe to modify the contents of the argument after Find returns.
func (p *DB) Find(key []byte) (rkey, value []byte, err error) {
	p.mu.RLock()
	if node, _ := p.findGE(key, false); node != 0 {
		n := p.nodeData[node]
		m := n + p.nodeData[node+nKey]
		rkey = p.kvData[n:m]
		value = p.kvData[m : m+p.nodeData[node+nVal]]
	} else {
		err = ErrNotFound
	}
	p.mu.RUnlock()
	return
}

// NewIterator returns an iterator of the DB.
// The returned iterator is not goroutine-safe, but it is safe to use
// multiple iterators concurrently, with each in a dedicated goroutine.
// It is also safe to use an iterator concurrently with modifying its
// underlying DB. However, the resultant key/value pairs are not guaranteed
// to be a consistent snapshot of the DB at a particular point in time.
//
// Slice allows slicing the iterator to only contains keys in the given
// range. A nil Range.Start is treated as a key before all keys in the
// DB. And a nil Range.Limit is treated as a key after all keys in
// the DB.
//
// The iterator must be released after use, by calling Release method.
//
// Also read Iterator documentation of the leveldb/iterator package.
func (p *DB) NewIterator(slice *util.Range) iterator.Iterator {
	return &dbIter{p: p, slice: slice}
}

// Capacity returns keys/values buffer capacity.
func (p *DB) Capacity() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return cap(p.kvData)
}

// Size returns sum of keys and values length. Note that deleted
// key/value will not be accouted for, but it will still consume
// the buffer, since the buffer is append only.
func (p *DB) Size() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.kvSize
}

// Free returns keys/values free buffer before need to grow.
func (p *DB) Free() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return cap(p.kvData) - len(p.kvData)
}

// Len returns the number of entries in the DB.
func (p *DB) Len() int {
	p.mu.RLock()
	defer p.mu.RUnlock()
	return p.n
}

// Reset resets the DB to initial empty state. Allows reuse the buffer.
func (p *DB) Reset() {
	p.rnd = rand.New(rand.NewSource(0xdeadbeef))
	p.maxHeight = 1
	p.n = 0
	p.kvSize = 0
	p.kvData = p.kvData[:0]
	p.nodeData = p.nodeData[:4+tMaxHeight]
	p.nodeData[nKV] = 0
	p.nodeData[nKey] = 0
	p.nodeData[nVal] = 0
	p.nodeData[nHeight] = tMaxHeight
	for n := 0; n < tMaxHeight; n++ {
		p.nodeData[4+n] = 0
		p.prevNode[n] = 0
	}
}

// New creates a new initalized in-memory key/value DB. The capacity
// is the initial key/value buffer capacity. The capacity is advisory,
// not enforced.
func New(cmp comparer.BasicComparer, capacity int) *DB {
	p := &DB{
		cmp:       cmp,
		rnd:       rand.New(rand.NewSource(0xdeadbeef)),
		maxHeight: 1,
		kvData:    make([]byte, 0, capacity),
		nodeData:  make([]int, 4+tMaxHeight),
	}
	p.nodeData[nHeight] = tMaxHeight
	return p
}
<|endoftext|>