<commit_before>package tester

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"

	"github.com/PuerkitoBio/goquery"
)

type Tester struct {
	mux *http.ServeMux
	t   *testing.T
}

func New(mux *http.ServeMux, t *testing.T) *Tester {
	return &Tester{
		mux: mux,
		t:   t,
	}
}

func (t *Tester) GetHTML(u string) (*goquery.Document, *httptest.ResponseRecorder) {
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		t.t.Fatal(err)
	}

	rw := httptest.NewRecorder()
	t.mux.ServeHTTP(rw, req)

	doc, err := goquery.NewDocumentFromReader(rw.Body)
	if err != nil {
		t.t.Fatal(err)
	}
	return doc, rw
}

func (t *Tester) GetBody(u string, h *http.Header) (rw *httptest.ResponseRecorder) {
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		t.t.Fatal(err)
	}

	if h != nil {
		req.Header = *h
	}

	rw = httptest.NewRecorder()
	t.mux.ServeHTTP(rw, req)
	return
}

func (t *Tester) GetJSON(u string, v interface{}) (rw *httptest.ResponseRecorder) {
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		t.t.Fatal(err)
	}

	rw = httptest.NewRecorder()
	t.mux.ServeHTTP(rw, req)
	body := rw.Body.Bytes()
	if rw.Code == 200 {
		if err = json.Unmarshal(body, v); err != nil {
			t.t.Log("Body:", string(body))
			t.t.Fatal(err)
		}
	}

	return
}

func (t *Tester) SendJSON(u string, v interface{}) (rw *httptest.ResponseRecorder) {
	blob, err := json.Marshal(v)
	if err != nil {
		t.t.Fatal(err)
	}

	req, err := http.NewRequest("POST", u, bytes.NewReader(blob))
	if err != nil {
		t.t.Fatal(err)
	}

	req.Header.Set("Content-Type", "application/json")

	rw = httptest.NewRecorder()
	t.mux.ServeHTTP(rw, req)
	return
}

func (t *Tester) PostForm(u string, v url.Values) (rw *httptest.ResponseRecorder) {
	req, err := http.NewRequest("POST", u, strings.NewReader(v.Encode()))
	if err != nil {
		t.t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	rw = httptest.NewRecorder()
	t.mux.ServeHTTP(rw, req)
	return
}
<commit_msg>added header argument to GetHTML<commit_after>package tester

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"

	"github.com/PuerkitoBio/goquery"
)

type Tester struct {
	mux *http.ServeMux
	t   *testing.T
}

func New(mux *http.ServeMux, t *testing.T) *Tester {
	return &Tester{
		mux: mux,
		t:   t,
	}
}

func (t *Tester) GetHTML(u string, h *http.Header) (*goquery.Document, *httptest.ResponseRecorder) {
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		t.t.Fatal(err)
	}

	if h != nil {
		req.Header = *h
	}

	rw := httptest.NewRecorder()
	t.mux.ServeHTTP(rw, req)

	doc, err := goquery.NewDocumentFromReader(rw.Body)
	if err != nil {
		t.t.Fatal(err)
	}
	return doc, rw
}

func (t *Tester) GetBody(u string, h *http.Header) (rw *httptest.ResponseRecorder) {
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		t.t.Fatal(err)
	}

	if h != nil {
		req.Header = *h
	}

	rw = httptest.NewRecorder()
	t.mux.ServeHTTP(rw, req)
	return
}

func (t *Tester) GetJSON(u string, v interface{}) (rw *httptest.ResponseRecorder) {
	req, err := http.NewRequest("GET", u, nil)
	if err != nil {
		t.t.Fatal(err)
	}

	rw = httptest.NewRecorder()
	t.mux.ServeHTTP(rw, req)
	body := rw.Body.Bytes()
	if rw.Code == 200 {
		if err = json.Unmarshal(body, v); err != nil {
			t.t.Log("Body:", string(body))
			t.t.Fatal(err)
		}
	}

	return
}

func (t *Tester) SendJSON(u string, v interface{}) (rw *httptest.ResponseRecorder) {
	blob, err := json.Marshal(v)
	if err != nil {
		t.t.Fatal(err)
	}

	req, err := http.NewRequest("POST", u, bytes.NewReader(blob))
	if err != nil {
		t.t.Fatal(err)
	}

	req.Header.Set("Content-Type", "application/json")

	rw = httptest.NewRecorder()
	t.mux.ServeHTTP(rw, req)
	return
}

func (t *Tester) PostForm(u string, v url.Values) (rw *httptest.ResponseRecorder) {
	req, err := http.NewRequest("POST", u, strings.NewReader(v.Encode()))
	if err != nil {
		t.t.Fatal(err)
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	rw = httptest.NewRecorder()
	t.mux.ServeHTTP(rw, req)
	return
}
<|endoftext|>


<commit_before>package httpevents

import (
	"net/http"

	"github.com/segmentio/events"
)

// NewHandler wraps the HTTP handler with and access logger which generates
// events for all incoming requests using logger.
//
// Panics from handler are intercepted and trigger a 500 response if no response
// header was sent yet. The panic is not slienced tho and is propagated to the
// parent handler.
func NewHandler(logger *events.Logger, handler http.Handler) http.Handler {
	if logger == nil {
		logger = events.DefaultLogger
	}
	return http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		w := &responseWriter{
			ResponseWriter: res,
			// We capture all the values we need from req in case the object
			// gets modified by the handler.
			logger:   logger,
			address:  req.RemoteAddr,
			method:   req.Method,
			host:     req.Host,
			path:     req.URL.Path,
			query:    req.URL.RawQuery,
			fragment: req.URL.Fragment,
			agent:    req.UserAgent(),
		}

		// If the handler panics we want to make sure we report the issue in the
		// access log, while also ensuring that a response is going to be sent
		// down to the client.
		// We don't silence the panic here tho and instead we forward it back to
		// the parent handler which may need to be aware that a panic occurred.
		defer func() {
			if err := recover(); err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				panic(err)
			}
		}()

		// The request is forwarded to the handler, if it never calls the
		// writer's WriteHeader method we force the call with "200 OK" status
		// to match the default behavior of the net/http package (and also make
		// sure an access log will be written).
		handler.ServeHTTP(w, req)
		w.WriteHeader(http.StatusOK)
	})
}

type responseWriter struct {
	http.ResponseWriter
	logger   *events.Logger
	address  string
	method   string
	host     string
	path     string
	query    string
	fragment string
	agent    string
}

func (w *responseWriter) WriteHeader(status int) {
	if w.logger != nil {
		var a [128]byte
		var b = append(a[:0], "%{address}s - %{host}s - %{method}s %{path}s"...)

		// Don't output a '?' character when the query string is empty, this is
		// a more natural way of reading URLs.
		if len(w.query) != 0 {
			b = append(b, '?')
		}
		b = append(b, "%{query}s"...)

		// Same than with the query string, don't output a '#' character when
		// there is no fragment.
		if len(w.fragment) != 0 {
			b = append(b, '#')
		}
		b = append(b, "%{fragment}s - %{status}d %s - %{agent}q"...)

		w.logger.Log(string(b),
			w.address,
			w.host,
			w.method,
			w.path,
			w.query,
			w.fragment,
			status,
			http.StatusText(status),
			w.agent,
		)
		w.logger = nil
	}
	w.ResponseWriter.WriteHeader(status)
}
<commit_msg>more meaningful variable names<commit_after>package httpevents

import (
	"net/http"

	"github.com/segmentio/events"
)

// NewHandler wraps the HTTP handler with and access logger which generates
// events for all incoming requests using logger.
//
// Panics from handler are intercepted and trigger a 500 response if no response
// header was sent yet. The panic is not slienced tho and is propagated to the
// parent handler.
func NewHandler(logger *events.Logger, handler http.Handler) http.Handler {
	if logger == nil {
		logger = events.DefaultLogger
	}
	return http.HandlerFunc(func(res http.ResponseWriter, req *http.Request) {
		w := &responseWriter{
			ResponseWriter: res,
			// We capture all the values we need from req in case the object
			// gets modified by the handler.
			logger:   logger,
			address:  req.RemoteAddr,
			method:   req.Method,
			host:     req.Host,
			path:     req.URL.Path,
			query:    req.URL.RawQuery,
			fragment: req.URL.Fragment,
			agent:    req.UserAgent(),
		}

		// If the handler panics we want to make sure we report the issue in the
		// access log, while also ensuring that a response is going to be sent
		// down to the client.
		// We don't silence the panic here tho and instead we forward it back to
		// the parent handler which may need to be aware that a panic occurred.
		defer func() {
			if err := recover(); err != nil {
				w.WriteHeader(http.StatusInternalServerError)
				panic(err)
			}
		}()

		// The request is forwarded to the handler, if it never calls the
		// writer's WriteHeader method we force the call with "200 OK" status
		// to match the default behavior of the net/http package (and also make
		// sure an access log will be written).
		handler.ServeHTTP(w, req)
		w.WriteHeader(http.StatusOK)
	})
}

type responseWriter struct {
	http.ResponseWriter
	logger   *events.Logger
	address  string
	method   string
	host     string
	path     string
	query    string
	fragment string
	agent    string
}

func (w *responseWriter) WriteHeader(status int) {
	if w.logger != nil {
		var buf [128]byte
		var fmt = append(buf[:0], "%{address}s - %{host}s - %{method}s %{path}s"...)

		// Don't output a '?' character when the query string is empty, this is
		// a more natural way of reading URLs.
		if len(w.query) != 0 {
			fmt = append(fmt, '?')
		}
		fmt = append(fmt, "%{query}s"...)

		// Same than with the query string, don't output a '#' character when
		// there is no fragment.
		if len(w.fragment) != 0 {
			fmt = append(fmt, '#')
		}
		fmt = append(fmt, "%{fragment}s - %{status}d %s - %{agent}q"...)

		w.logger.Log(string(fmt),
			w.address,
			w.host,
			w.method,
			w.path,
			w.query,
			w.fragment,
			status,
			http.StatusText(status),
			w.agent,
		)
		w.logger = nil
	}
	w.ResponseWriter.WriteHeader(status)
}
<|endoftext|>


<commit_before>package httpservice

import (
	"net/http"

	"code.google.com/p/go.net/context"

	"github.com/arjantop/saola"
	"github.com/julienschmidt/httprouter"
)

type key int

const httpRequest key = 0

type HttpRequest struct {
	Writer  http.ResponseWriter
	Request *http.Request
}

func WithHttpRequest(ctx context.Context, w http.ResponseWriter, r *http.Request) context.Context {
	return context.WithValue(ctx, httpRequest, &HttpRequest{w, r})
}

func GetHttpRequest(ctx context.Context) *HttpRequest {
	r, ok := ctx.Value(httpRequest).(*HttpRequest)
	if !ok {
		panic("missing http request in context")
	}
	return r
}

type Params struct {
	params httprouter.Params
}

func EmptyParams() Params {
	return Params{make([]httprouter.Param, 0)}
}

func (p Params) Get(key string) string {
	return p.params.ByName(key)
}

func (p *Params) Set(key, value string) {
	for i, _ := range p.params {
		p := &p.params[i]
		if p.Key == key {
			p.Value = value
			return
		}
	}
	p.params = append(p.params, httprouter.Param{
		Key:   key,
		Value: value,
	})
}

const paramsKey key = 1

func WithParams(ctx context.Context, p Params) context.Context {
	return context.WithValue(ctx, paramsKey, p)
}

func GetParams(ctx context.Context) Params {
	if params, ok := ctx.Value(paramsKey).(Params); ok {
		return params
	}
	return EmptyParams()
}

type HttpService interface {
	DoHTTP(ctx context.Context, w http.ResponseWriter, r *http.Request) error
	saola.Service
}

type FuncService func(ctx context.Context, w http.ResponseWriter, r *http.Request) error

func (f FuncService) DoHTTP(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	return f(ctx, w, r)
}

func (f FuncService) Do(ctx context.Context) error {
	r := GetHttpRequest(ctx)
	return f.DoHTTP(ctx, r.Writer, r.Request)
}

type Endpoint struct {
	router *httprouter.Router
}

func NewEndpoint() *Endpoint {
	return &Endpoint{
		router: httprouter.New(),
	}
}

func (e *Endpoint) GET(path string, s saola.Service) {
	e.router.GET(path, func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		ctx := WithParams(WithHttpRequest(context.Background(), w, r), Params{p})
		s.Do(ctx)
	})
}

func (e *Endpoint) POST(path string, s saola.Service) {
	e.router.POST(path, func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		ctx := WithParams(WithHttpRequest(context.Background(), w, r), Params{p})
		s.Do(ctx)
	})
}

func (e *Endpoint) PUT(path string, s saola.Service) {
	e.router.PUT(path, func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		ctx := WithParams(WithHttpRequest(context.Background(), w, r), Params{p})
		s.Do(ctx)
	})
}

func (e *Endpoint) DoHTTP(_ context.Context, w http.ResponseWriter, r *http.Request) error {
	e.router.ServeHTTP(w, r)
	return nil
}

func (e *Endpoint) Do(ctx context.Context) error {
	r := GetHttpRequest(ctx)
	return e.DoHTTP(ctx, r.Writer, r.Request)
}

func Serve(addr string, s saola.Service) error {
	return http.ListenAndServe(addr, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := WithHttpRequest(context.Background(), w, r)
		s.Do(ctx)
	}))
}
<commit_msg>Utility method for implementing HttpService.<commit_after>package httpservice

import (
	"log"
	"net/http"

	"code.google.com/p/go.net/context"

	"github.com/arjantop/saola"
	"github.com/julienschmidt/httprouter"
)

type key int

const httpRequest key = 0

type HttpRequest struct {
	Writer  http.ResponseWriter
	Request *http.Request
}

func WithHttpRequest(ctx context.Context, w http.ResponseWriter, r *http.Request) context.Context {
	return context.WithValue(ctx, httpRequest, &HttpRequest{w, r})
}

func GetHttpRequest(ctx context.Context) *HttpRequest {
	r, ok := ctx.Value(httpRequest).(*HttpRequest)
	if !ok {
		panic("missing http request in context")
	}
	return r
}

type Params struct {
	params httprouter.Params
}

func EmptyParams() Params {
	return Params{make([]httprouter.Param, 0)}
}

func (p Params) Get(key string) string {
	return p.params.ByName(key)
}

func (p *Params) Set(key, value string) {
	for i, _ := range p.params {
		p := &p.params[i]
		if p.Key == key {
			p.Value = value
			return
		}
	}
	p.params = append(p.params, httprouter.Param{
		Key:   key,
		Value: value,
	})
}

const paramsKey key = 1

func WithParams(ctx context.Context, p Params) context.Context {
	return context.WithValue(ctx, paramsKey, p)
}

func GetParams(ctx context.Context) Params {
	if params, ok := ctx.Value(paramsKey).(Params); ok {
		return params
	}
	return EmptyParams()
}

type HttpService interface {
	DoHTTP(ctx context.Context, w http.ResponseWriter, r *http.Request) error
	saola.Service
}

type FuncService func(ctx context.Context, w http.ResponseWriter, r *http.Request) error

func (f FuncService) DoHTTP(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
	return f(ctx, w, r)
}

func (f FuncService) Do(ctx context.Context) error {
	r := GetHttpRequest(ctx)
	return f.DoHTTP(ctx, r.Writer, r.Request)
}

type Endpoint struct {
	router *httprouter.Router
}

func NewEndpoint() *Endpoint {
	log.Println("")
	return &Endpoint{
		router: httprouter.New(),
	}
}

func (e *Endpoint) GET(path string, s saola.Service) {
	e.router.GET(path, func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		ctx := WithParams(WithHttpRequest(context.Background(), w, r), Params{p})
		s.Do(ctx)
	})
}

func (e *Endpoint) POST(path string, s saola.Service) {
	e.router.POST(path, func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		ctx := WithParams(WithHttpRequest(context.Background(), w, r), Params{p})
		s.Do(ctx)
	})
}

func (e *Endpoint) PUT(path string, s saola.Service) {
	e.router.PUT(path, func(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
		ctx := WithParams(WithHttpRequest(context.Background(), w, r), Params{p})
		s.Do(ctx)
	})
}

func (e *Endpoint) DoHTTP(_ context.Context, w http.ResponseWriter, r *http.Request) error {
	e.router.ServeHTTP(w, r)
	return nil
}

func (e *Endpoint) Do(ctx context.Context) error {
	return Do(e, ctx)
}

func Serve(addr string, s saola.Service) error {
	return http.ListenAndServe(addr, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := WithHttpRequest(context.Background(), w, r)
		s.Do(ctx)
	}))
}

func Do(s HttpService, ctx context.Context) error {
	r := GetHttpRequest(ctx)
	return s.DoHTTP(ctx, r.Writer, r.Request)
}
<|endoftext|>


<commit_before>// Copyright 2013 Hui Chen
// Copyright 2016 ego 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 riot

import (
	// "fmt"

	"strings"

	"github.com/go-ego/gpy"
	"github.com/go-ego/riot/types"
)

type segmenterRequest struct {
	docId uint64
	hash  uint32
	data  types.DocIndexData
	// data        types.DocumentIndexData
	forceUpdate bool
}

// Map defines the type map[string][]int
type Map map[string][]int

// Segspl split seg
func (engine *Engine) Segspl(splData []string, num int) (Map, int) {
	var (
		numTokens int
		sqlitStr  string
	)
	tokensMap := make(map[string][]int)
	for i := 0; i < num; i++ {
		if splData[i] != "" {
			if !engine.stopTokens.IsStopToken(splData[i]) {
				numTokens++
				tokensMap[splData[i]] = append(tokensMap[splData[i]], numTokens)
			}

			sqlitStr += splData[i]
			if !engine.stopTokens.IsStopToken(sqlitStr) {
				numTokens++
				tokensMap[sqlitStr] = append(tokensMap[sqlitStr], numTokens)
			}

			if engine.initOptions.Using == 6 {
				// more combination
				var sqlitsStr string
				for s := i + 1; s < len(splData); s++ {
					sqlitsStr += splData[s]

					if !engine.stopTokens.IsStopToken(sqlitsStr) {
						numTokens++
						tokensMap[sqlitsStr] = append(tokensMap[sqlitsStr], numTokens)
					}
				}
			}

		}
	}

	return tokensMap, numTokens
}

func (engine *Engine) splitData(request segmenterRequest) (Map, int) {
	tokensMap := make(map[string][]int)
	var (
		num       int
		numTokens int
	)

	if request.data.Content != "" {
		request.data.Content = strings.ToLower(request.data.Content)
		if engine.initOptions.Using == 3 {
			// use segmenter
			segments := engine.segmenter.Segment([]byte(request.data.Content))
			for _, segment := range segments {
				token := segment.Token().Text()
				if !engine.stopTokens.IsStopToken(token) {
					tokensMap[token] = append(tokensMap[token], segment.Start())
				}
			}
			numTokens += len(segments)
		}

		if engine.initOptions.Using == 4 {
			// use segmenter
			splSpaData := strings.Split(request.data.Content, " ")
			num := len(splSpaData)
			tokenMap, numToken := engine.Segspl(splSpaData, num)
			numTokens += numToken
			for key, val := range tokenMap {
				tokensMap[key] = val
			}
		}

		if engine.initOptions.Using != 4 {
			splData := strings.Split(request.data.Content, "")
			num = len(splData)
			tokenMap, numToken := engine.Segspl(splData, num)
			numTokens += numToken
			for key, val := range tokenMap {
				tokensMap[key] = val
			}
		}
	}

	for _, t := range request.data.Tokens {
		if !engine.stopTokens.IsStopToken(t.Text) {
			tokensMap[t.Text] = t.Locations
		}
	}

	numTokens += len(request.data.Tokens)

	return tokensMap, numTokens
}

func (engine *Engine) segmenterData(request segmenterRequest) (Map, int) {
	tokensMap := make(map[string][]int)
	numTokens := 0

	if engine.initOptions.Using == 0 && request.data.Content != "" {
		// Content分词, 当文档正文不为空时，优先从内容分词中得到关键词
		segments := engine.segmenter.Segment([]byte(request.data.Content))
		for _, segment := range segments {
			token := segment.Token().Text()
			if !engine.stopTokens.IsStopToken(token) {
				tokensMap[token] = append(tokensMap[token], segment.Start())
			}
		}

		for _, t := range request.data.Tokens {
			if !engine.stopTokens.IsStopToken(t.Text) {
				tokensMap[t.Text] = t.Locations
			}
		}

		numTokens = len(segments) + len(request.data.Tokens)

		return tokensMap, numTokens
	}

	if engine.initOptions.Using == 1 && request.data.Content != "" {
		// Content分词, 当文档正文不为空时，优先从内容分词中得到关键词
		segments := engine.segmenter.Segment([]byte(request.data.Content))
		for _, segment := range segments {
			token := segment.Token().Text()
			if !engine.stopTokens.IsStopToken(token) {
				tokensMap[token] = append(tokensMap[token], segment.Start())
			}
		}
		numTokens = len(segments)

		return tokensMap, numTokens
	}

	if engine.initOptions.Using == 2 ||
		((engine.initOptions.Using == 1 || engine.initOptions.Using == 3) &&
			request.data.Content == "") {
		for _, t := range request.data.Tokens {
			if !engine.stopTokens.IsStopToken(t.Text) {
				tokensMap[t.Text] = t.Locations
			}
		}

		numTokens = len(request.data.Tokens)

		return tokensMap, numTokens
	}

	tokenMap, lenSplitData := engine.splitData(request)

	return tokenMap, lenSplitData
}

func (engine *Engine) segmenterWorker() {
	for {
		request := <-engine.segmenterChannel
		if request.docId == 0 {
			if request.forceUpdate {
				for i := 0; i < engine.initOptions.NumShards; i++ {
					engine.indexerAddDocChannels[i] <- indexerAddDocumentRequest{forceUpdate: true}
				}
			}
			continue
		}

		shard := engine.getShard(request.hash)
		tokensMap, numTokens := engine.segmenterData(request)

		// 加入非分词的文档标签
		for _, label := range request.data.Labels {
			if !engine.initOptions.NotUsingSegmenter {
				if !engine.stopTokens.IsStopToken(label) {
					//当正文中已存在关键字时，若不判断，位置信息将会丢失
					if _, ok := tokensMap[label]; !ok {
						tokensMap[label] = []int{}
					}
				}
			} else {
				//当正文中已存在关键字时，若不判断，位置信息将会丢失
				if _, ok := tokensMap[label]; !ok {
					tokensMap[label] = []int{}
				}
			}
		}

		indexerRequest := indexerAddDocumentRequest{
			document: &types.DocumentIndex{
				DocId:       request.docId,
				TokenLength: float32(numTokens),
				Keywords:    make([]types.KeywordIndex, len(tokensMap)),
			},
			forceUpdate: request.forceUpdate,
		}
		iTokens := 0
		for k, v := range tokensMap {
			indexerRequest.document.Keywords[iTokens] = types.KeywordIndex{
				Text: k,
				// 非分词标注的词频设置为0，不参与tf-idf计算
				Frequency: float32(len(v)),
				Starts:    v}
			iTokens++
		}

		engine.indexerAddDocChannels[shard] <- indexerRequest
		if request.forceUpdate {
			for i := 0; i < engine.initOptions.NumShards; i++ {
				if i == shard {
					continue
				}
				engine.indexerAddDocChannels[i] <- indexerAddDocumentRequest{forceUpdate: true}
			}
		}
		rankerRequest := rankerAddDocRequest{
			// docId: request.docId, fields: request.data.Fields}
			docId: request.docId, fields: request.data.Fields, content: request.data.Content, attri: request.data.Attri}
		engine.rankerAddDocChannels[shard] <- rankerRequest
	}
}

// PinYin get the Chinese alphabet and abbreviation
func (engine *Engine) PinYin(hans string) []string {
	var (
		str      string
		pystr    string
		strArr   []string
		splitStr string
		// splitArr []string
	)

	//
	splitHans := strings.Split(hans, "")
	for i := 0; i < len(splitHans); i++ {
		if splitHans[i] != "" {
			if !engine.stopTokens.IsStopToken(splitHans[i]) {
				strArr = append(strArr, splitHans[i])
			}
			splitStr += splitHans[i]
		}
		if !engine.stopTokens.IsStopToken(splitStr) {
			strArr = append(strArr, splitStr)
		}
	}

	// Segment 分词
	if !engine.initOptions.NotUsingSegmenter {
		sehans := engine.Segment(hans)
		for h := 0; h < len(sehans); h++ {
			if !engine.stopTokens.IsStopToken(sehans[h]) {
				strArr = append(strArr, sehans[h])
			}
		}
	}
	//
	// py := pinyin.LazyConvert(sehans[h], nil)
	py := gpy.LazyConvert(hans, nil)

	// log.Println("py...", py)
	for i := 0; i < len(py); i++ {
		// log.Println("py[i]...", py[i])
		pystr += py[i]
		if !engine.stopTokens.IsStopToken(pystr) {
			strArr = append(strArr, pystr)
		}

		if len(py[i]) > 0 {
			str += py[i][0:1]
			if !engine.stopTokens.IsStopToken(str) {
				strArr = append(strArr, str)
			}
		}
	}

	return strArr
}
<commit_msg>fix spelling error and update code<commit_after>// Copyright 2013 Hui Chen
// Copyright 2016 ego 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 riot

import (
	// "fmt"

	"strings"

	"github.com/go-ego/gpy"
	"github.com/go-ego/riot/types"
)

type segmenterRequest struct {
	docId uint64
	hash  uint32
	data  types.DocIndexData
	// data        types.DocumentIndexData
	forceUpdate bool
}

// Map defines the type map[string][]int
type Map map[string][]int

// ForSplitData for split seg data, segspl
func (engine *Engine) ForSplitData(splData []string, num int) (Map, int) {
	var (
		numTokens int
		splitStr  string
	)
	tokensMap := make(map[string][]int)
	for i := 0; i < num; i++ {
		if splData[i] != "" {
			if !engine.stopTokens.IsStopToken(splData[i]) {
				numTokens++
				tokensMap[splData[i]] = append(tokensMap[splData[i]], numTokens)
			}

			splitStr += splData[i]
			if !engine.stopTokens.IsStopToken(splitStr) {
				numTokens++
				tokensMap[splitStr] = append(tokensMap[splitStr], numTokens)
			}

			if engine.initOptions.Using == 6 {
				// more combination
				var splitsStr string
				for s := i + 1; s < len(splData); s++ {
					splitsStr += splData[s]

					if !engine.stopTokens.IsStopToken(splitsStr) {
						numTokens++
						tokensMap[splitsStr] = append(tokensMap[splitsStr], numTokens)
					}
				}
			}

		}
	}

	return tokensMap, numTokens
}

func (engine *Engine) splitData(request segmenterRequest) (Map, int) {
	tokensMap := make(map[string][]int)
	var (
		num       int
		numTokens int
	)

	if request.data.Content != "" {
		request.data.Content = strings.ToLower(request.data.Content)
		if engine.initOptions.Using == 3 {
			// use segmenter
			segments := engine.segmenter.Segment([]byte(request.data.Content))
			for _, segment := range segments {
				token := segment.Token().Text()
				if !engine.stopTokens.IsStopToken(token) {
					tokensMap[token] = append(tokensMap[token], segment.Start())
				}
			}
			numTokens += len(segments)
		}

		if engine.initOptions.Using == 4 {
			// use segmenter
			splSpaData := strings.Split(request.data.Content, " ")
			num := len(splSpaData)
			tokenMap, numToken := engine.ForSplitData(splSpaData, num)
			numTokens += numToken
			for key, val := range tokenMap {
				tokensMap[key] = val
			}
		}

		if engine.initOptions.Using != 4 {
			splData := strings.Split(request.data.Content, "")
			num = len(splData)
			tokenMap, numToken := engine.ForSplitData(splData, num)
			numTokens += numToken
			for key, val := range tokenMap {
				tokensMap[key] = val
			}
		}
	}

	for _, t := range request.data.Tokens {
		if !engine.stopTokens.IsStopToken(t.Text) {
			tokensMap[t.Text] = t.Locations
		}
	}

	numTokens += len(request.data.Tokens)

	return tokensMap, numTokens
}

func (engine *Engine) segmenterData(request segmenterRequest) (Map, int) {
	tokensMap := make(map[string][]int)
	numTokens := 0

	if engine.initOptions.Using == 0 && request.data.Content != "" {
		// Content分词, 当文档正文不为空时，优先从内容分词中得到关键词
		segments := engine.segmenter.Segment([]byte(request.data.Content))
		for _, segment := range segments {
			token := segment.Token().Text()
			if !engine.stopTokens.IsStopToken(token) {
				tokensMap[token] = append(tokensMap[token], segment.Start())
			}
		}

		for _, t := range request.data.Tokens {
			if !engine.stopTokens.IsStopToken(t.Text) {
				tokensMap[t.Text] = t.Locations
			}
		}

		numTokens = len(segments) + len(request.data.Tokens)

		return tokensMap, numTokens
	}

	if engine.initOptions.Using == 1 && request.data.Content != "" {
		// Content分词, 当文档正文不为空时，优先从内容分词中得到关键词
		segments := engine.segmenter.Segment([]byte(request.data.Content))
		for _, segment := range segments {
			token := segment.Token().Text()
			if !engine.stopTokens.IsStopToken(token) {
				tokensMap[token] = append(tokensMap[token], segment.Start())
			}
		}
		numTokens = len(segments)

		return tokensMap, numTokens
	}

	if engine.initOptions.Using == 2 ||
		((engine.initOptions.Using == 1 || engine.initOptions.Using == 3) &&
			request.data.Content == "") {
		for _, t := range request.data.Tokens {
			if !engine.stopTokens.IsStopToken(t.Text) {
				tokensMap[t.Text] = t.Locations
			}
		}

		numTokens = len(request.data.Tokens)

		return tokensMap, numTokens
	}

	tokenMap, lenSplitData := engine.splitData(request)

	return tokenMap, lenSplitData
}

func (engine *Engine) segmenterWorker() {
	for {
		request := <-engine.segmenterChannel
		if request.docId == 0 {
			if request.forceUpdate {
				for i := 0; i < engine.initOptions.NumShards; i++ {
					engine.indexerAddDocChannels[i] <- indexerAddDocumentRequest{forceUpdate: true}
				}
			}
			continue
		}

		shard := engine.getShard(request.hash)
		tokensMap, numTokens := engine.segmenterData(request)

		// 加入非分词的文档标签
		for _, label := range request.data.Labels {
			if !engine.initOptions.NotUsingSegmenter {
				if !engine.stopTokens.IsStopToken(label) {
					//当正文中已存在关键字时，若不判断，位置信息将会丢失
					if _, ok := tokensMap[label]; !ok {
						tokensMap[label] = []int{}
					}
				}
			} else {
				//当正文中已存在关键字时，若不判断，位置信息将会丢失
				if _, ok := tokensMap[label]; !ok {
					tokensMap[label] = []int{}
				}
			}
		}

		indexerRequest := indexerAddDocumentRequest{
			document: &types.DocumentIndex{
				DocId:       request.docId,
				TokenLength: float32(numTokens),
				Keywords:    make([]types.KeywordIndex, len(tokensMap)),
			},
			forceUpdate: request.forceUpdate,
		}
		iTokens := 0
		for k, v := range tokensMap {
			indexerRequest.document.Keywords[iTokens] = types.KeywordIndex{
				Text: k,
				// 非分词标注的词频设置为0，不参与tf-idf计算
				Frequency: float32(len(v)),
				Starts:    v}
			iTokens++
		}

		engine.indexerAddDocChannels[shard] <- indexerRequest
		if request.forceUpdate {
			for i := 0; i < engine.initOptions.NumShards; i++ {
				if i == shard {
					continue
				}
				engine.indexerAddDocChannels[i] <- indexerAddDocumentRequest{forceUpdate: true}
			}
		}
		rankerRequest := rankerAddDocRequest{
			// docId: request.docId, fields: request.data.Fields}
			docId: request.docId, fields: request.data.Fields, content: request.data.Content, attri: request.data.Attri}
		engine.rankerAddDocChannels[shard] <- rankerRequest
	}
}

// PinYin get the Chinese alphabet and abbreviation
func (engine *Engine) PinYin(hans string) []string {
	var (
		str      string
		pystr    string
		strArr   []string
		splitStr string
		// splitArr []string
	)

	//
	splitHans := strings.Split(hans, "")
	for i := 0; i < len(splitHans); i++ {
		if splitHans[i] != "" {
			if !engine.stopTokens.IsStopToken(splitHans[i]) {
				strArr = append(strArr, splitHans[i])
			}
			splitStr += splitHans[i]
		}
		if !engine.stopTokens.IsStopToken(splitStr) {
			strArr = append(strArr, splitStr)
		}
	}

	// Segment 分词
	if !engine.initOptions.NotUsingSegmenter {
		sehans := engine.Segment(hans)
		for h := 0; h < len(sehans); h++ {
			if !engine.stopTokens.IsStopToken(sehans[h]) {
				strArr = append(strArr, sehans[h])
			}
		}
	}
	//
	// py := pinyin.LazyConvert(sehans[h], nil)
	py := gpy.LazyConvert(hans, nil)

	// log.Println("py...", py)
	for i := 0; i < len(py); i++ {
		// log.Println("py[i]...", py[i])
		pystr += py[i]
		if !engine.stopTokens.IsStopToken(pystr) {
			strArr = append(strArr, pystr)
		}

		if len(py[i]) > 0 {
			str += py[i][0:1]
			if !engine.stopTokens.IsStopToken(str) {
				strArr = append(strArr, str)
			}
		}
	}

	return strArr
}
<|endoftext|>


<commit_before>package sentinel

import (
	"github.com/mdevilliers/redishappy/configuration"
	"github.com/mdevilliers/redishappy/services/logger"
	"github.com/mdevilliers/redishappy/services/redis"
	"github.com/mdevilliers/redishappy/types"
	"github.com/mdevilliers/redishappy/util"
	"sync"
	"time"
)

const (
	SentinelMarkedUp    = 1
	SentinelMarkedDown  = 2
	SentinelMarkedAlive = 3
)

type Manager interface {
	Notify(event SentinelEvent)
	GetState(request TopologyRequest)
}

type SentinelManager struct {
	eventsChannel          chan SentinelEvent
	topologyRequestChannel chan TopologyRequest
	switchmasterchannel    chan types.MasterSwitchedEvent
	redisConnection        redis.RedisConnection
}

var topologyState = SentinelTopology{Sentinels: map[string]*SentinelInfo{}}
var statelock = &sync.Mutex{}

func NewManager(switchmasterchannel chan types.MasterSwitchedEvent, configuration *configuration.Configuration) *SentinelManager {

	events := make(chan SentinelEvent)
	requests := make(chan TopologyRequest)

	manager := &SentinelManager{eventsChannel: events,
		topologyRequestChannel: requests,
		switchmasterchannel:    switchmasterchannel,
		redisConnection:        redis.RadixRedisConnection{}}

	go loopEvents(events, requests, manager)

	for _, sentinel := range configuration.Sentinels {
		manager.Notify(&SentinelAdded{Sentinel: sentinel})
	}

	return manager
}

func (m *SentinelManager) Notify(event SentinelEvent) {
	m.eventsChannel <- event
}

func (m *SentinelManager) GetState(request TopologyRequest) {
	m.topologyRequestChannel <- request
}

func (m *SentinelManager) ClearState() {
	statelock.Lock()
	defer statelock.Unlock()
	topologyState = SentinelTopology{Sentinels: map[string]*SentinelInfo{}}
}

func (m *SentinelManager) GetTopology(stateChannel chan types.MasterDetailsCollection, configuration *configuration.Configuration) {

	topology := types.NewMasterDetailsCollection()

	for _, sentinel := range configuration.Sentinels {

		client, err := NewSentinelClient(sentinel, m.redisConnection)
		defer client.Close()

		if err != nil {
			logger.Info.Printf("Error starting sentinel (%s) client : %s", sentinel.GetLocation(), err.Error())
			continue
		}

		for _, clusterDetails := range configuration.Clusters {

			details, err := client.DiscoverMasterForCluster(clusterDetails.Name)

			if err != nil {
				continue
			}

			details.ExternalPort = clusterDetails.MasterPort
			// TODO : last one wins?
			topology.AddOrReplace(details)

			m.exploreSentinelTopology(client, clusterDetails.Name)
		}
	}
	stateChannel <- topology
}

func (m *SentinelManager) exploreSentinelTopology(client *SentinelClient, clustername string) {

	sentinels := client.FindConnectedSentinels(clustername)
	if len(sentinels) > 0 {
		m.notifySentinelsAreConnected(sentinels)
	}
}

func (m *SentinelManager) notifySentinelsAreConnected(sentinels []types.Sentinel) {
	for _, sentinel := range sentinels {
		m.Notify(&SentinelAdded{Sentinel: sentinel})
	}
}

func (m *SentinelManager) startNewMonitor(sentinel types.Sentinel) (*Monitor, error) {

	monitor, err := NewMonitor(sentinel, m, m.redisConnection)

	if err != nil {
		logger.Error.Printf("Error starting monitor %s : %s", sentinel.GetLocation(), err.Error())
		m.Notify(&SentinelLost{Sentinel: sentinel})
		return nil, err
	}

	go monitor.StartMonitoringMasterEvents(m.switchmasterchannel)

	return monitor, nil
}

func loopEvents(events chan SentinelEvent, topology chan TopologyRequest, m *SentinelManager) {
	for {
		select {
		case event := <-events:
			updateState(event, m)
		case read := <-topology:
			read.ReplyChannel <- topologyState
		}
	}
}

func updateState(event interface{}, m *SentinelManager) {

	statelock.Lock()
	defer statelock.Unlock()

	switch e := event.(type) {
	case *SentinelAdded:

		sentinel := e.GetSentinel()
		uid := topologyState.createKey(sentinel)

		//if we don't know about the sentinel start monitoring it
		if _, exists := topologyState.Sentinels[uid]; !exists {

			info := &SentinelInfo{SentinelLocation: uid,
				LastUpdated:   time.Now().UTC(),
				KnownClusters: []ClusterInfo{},
				State:         SentinelMarkedUp}

			topologyState.Sentinels[uid] = info

			go m.startNewMonitor(sentinel)

			logger.Trace.Printf("Sentinel added : %s", util.String(topologyState))
		}

	case *SentinelLost:

		sentinel := e.GetSentinel()
		uid := topologyState.createKey(sentinel)
		currentInfo, ok := topologyState.Sentinels[uid]

		if ok {
			currentInfo.State = SentinelMarkedDown
			currentInfo.LastUpdated = time.Now().UTC()
		}

		util.Schedule(func() { m.startNewMonitor(sentinel) }, time.Second*5)

		logger.Trace.Printf("Sentinel lost : %s (scheduling new client and monitor).", util.String(topologyState))

	case *SentinelPing:
		sentinel := e.GetSentinel()
		uid := topologyState.createKey(sentinel)
		currentInfo, ok := topologyState.Sentinels[uid]

		if ok {
			currentInfo.State = SentinelMarkedAlive
			currentInfo.LastUpdated = time.Now().UTC()
			// currentInfo.KnownClusters = e.Clusters
		}

	default:
		logger.Error.Println("Unknown sentinel event : ", util.String(e))
	}
}
<commit_msg>only close a client if succesfully opened a connection<commit_after>package sentinel

import (
	"github.com/mdevilliers/redishappy/configuration"
	"github.com/mdevilliers/redishappy/services/logger"
	"github.com/mdevilliers/redishappy/services/redis"
	"github.com/mdevilliers/redishappy/types"
	"github.com/mdevilliers/redishappy/util"
	"sync"
	"time"
)

const (
	SentinelMarkedUp    = 1
	SentinelMarkedDown  = 2
	SentinelMarkedAlive = 3
)

type Manager interface {
	Notify(event SentinelEvent)
	GetState(request TopologyRequest)
}

type SentinelManager struct {
	eventsChannel          chan SentinelEvent
	topologyRequestChannel chan TopologyRequest
	switchmasterchannel    chan types.MasterSwitchedEvent
	redisConnection        redis.RedisConnection
}

var topologyState = SentinelTopology{Sentinels: map[string]*SentinelInfo{}}
var statelock = &sync.Mutex{}

func NewManager(switchmasterchannel chan types.MasterSwitchedEvent, configuration *configuration.Configuration) *SentinelManager {

	events := make(chan SentinelEvent)
	requests := make(chan TopologyRequest)

	manager := &SentinelManager{eventsChannel: events,
		topologyRequestChannel: requests,
		switchmasterchannel:    switchmasterchannel,
		redisConnection:        redis.RadixRedisConnection{}}

	go loopEvents(events, requests, manager)

	for _, sentinel := range configuration.Sentinels {
		manager.Notify(&SentinelAdded{Sentinel: sentinel})
	}

	return manager
}

func (m *SentinelManager) Notify(event SentinelEvent) {
	m.eventsChannel <- event
}

func (m *SentinelManager) GetState(request TopologyRequest) {
	m.topologyRequestChannel <- request
}

func (m *SentinelManager) ClearState() {
	statelock.Lock()
	defer statelock.Unlock()
	topologyState = SentinelTopology{Sentinels: map[string]*SentinelInfo{}}
}

func (m *SentinelManager) GetTopology(stateChannel chan types.MasterDetailsCollection, configuration *configuration.Configuration) {

	topology := types.NewMasterDetailsCollection()

	for _, sentinel := range configuration.Sentinels {

		client, err := NewSentinelClient(sentinel, m.redisConnection)

		if err != nil {
			logger.Info.Printf("Error starting sentinel (%s) client : %s", sentinel.GetLocation(), err.Error())
			continue
		}
		defer client.Close()

		for _, clusterDetails := range configuration.Clusters {

			details, err := client.DiscoverMasterForCluster(clusterDetails.Name)

			if err != nil {
				continue
			}

			details.ExternalPort = clusterDetails.MasterPort
			// TODO : last one wins?
			topology.AddOrReplace(details)

			m.exploreSentinelTopology(client, clusterDetails.Name)
		}
	}
	stateChannel <- topology
}

func (m *SentinelManager) exploreSentinelTopology(client *SentinelClient, clustername string) {

	sentinels := client.FindConnectedSentinels(clustername)
	if len(sentinels) > 0 {
		m.notifySentinelsAreConnected(sentinels)
	}
}

func (m *SentinelManager) notifySentinelsAreConnected(sentinels []types.Sentinel) {
	for _, sentinel := range sentinels {
		m.Notify(&SentinelAdded{Sentinel: sentinel})
	}
}

func (m *SentinelManager) startNewMonitor(sentinel types.Sentinel) (*Monitor, error) {

	monitor, err := NewMonitor(sentinel, m, m.redisConnection)

	if err != nil {
		logger.Error.Printf("Error starting monitor %s : %s", sentinel.GetLocation(), err.Error())
		m.Notify(&SentinelLost{Sentinel: sentinel})
		return nil, err
	}

	go monitor.StartMonitoringMasterEvents(m.switchmasterchannel)

	return monitor, nil
}

func loopEvents(events chan SentinelEvent, topology chan TopologyRequest, m *SentinelManager) {
	for {
		select {
		case event := <-events:
			updateState(event, m)
		case read := <-topology:
			read.ReplyChannel <- topologyState
		}
	}
}

func updateState(event interface{}, m *SentinelManager) {

	statelock.Lock()
	defer statelock.Unlock()

	switch e := event.(type) {
	case *SentinelAdded:

		sentinel := e.GetSentinel()
		uid := topologyState.createKey(sentinel)

		//if we don't know about the sentinel start monitoring it
		if _, exists := topologyState.Sentinels[uid]; !exists {

			info := &SentinelInfo{SentinelLocation: uid,
				LastUpdated:   time.Now().UTC(),
				KnownClusters: []ClusterInfo{},
				State:         SentinelMarkedUp}

			topologyState.Sentinels[uid] = info

			go m.startNewMonitor(sentinel)

			logger.Trace.Printf("Sentinel added : %s", util.String(topologyState))
		}

	case *SentinelLost:

		sentinel := e.GetSentinel()
		uid := topologyState.createKey(sentinel)
		currentInfo, ok := topologyState.Sentinels[uid]

		if ok {
			currentInfo.State = SentinelMarkedDown
			currentInfo.LastUpdated = time.Now().UTC()
		}

		util.Schedule(func() { m.startNewMonitor(sentinel) }, time.Second*5)

		logger.Trace.Printf("Sentinel lost : %s (scheduling new client and monitor).", util.String(topologyState))

	case *SentinelPing:
		sentinel := e.GetSentinel()
		uid := topologyState.createKey(sentinel)
		currentInfo, ok := topologyState.Sentinels[uid]

		if ok {
			currentInfo.State = SentinelMarkedAlive
			currentInfo.LastUpdated = time.Now().UTC()
			// currentInfo.KnownClusters = e.Clusters
		}

	default:
		logger.Error.Println("Unknown sentinel event : ", util.String(e))
	}
}
<|endoftext|>


<commit_before>package trackerserver

import (
	"encoding/json"
	"fmt"
	"net/http"

	"code.uber.internal/infra/kraken/core"
	"code.uber.internal/infra/kraken/tracker/announceclient"
	"code.uber.internal/infra/kraken/tracker/storage"
	"code.uber.internal/infra/kraken/utils/errutil"
	"code.uber.internal/infra/kraken/utils/handler"
	"code.uber.internal/infra/kraken/utils/log"
)

func (s *Server) announceHandler(w http.ResponseWriter, r *http.Request) error {
	req := new(announceclient.Request)
	if err := json.NewDecoder(r.Body).Decode(req); err != nil {
		return handler.Errorf("json decode request: %s", err)
	}
	if err := s.peerStore.UpdatePeer(req.InfoHash, req.Peer); err != nil {
		log.With(
			"hash", req.InfoHash,
			"peer_id", req.Peer.PeerID).Errorf("Error updating peer: %s", err)
	}
	peers, err := s.getPeerHandout(req.Name, req.InfoHash, req.Peer)
	if err != nil {
		return err
	}
	resp := &announceclient.Response{
		Peers:    peers,
		Interval: s.config.AnnounceInterval,
	}
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		return handler.Errorf("json encode response: %s", err)
	}
	return nil
}

func (s *Server) getPeerHandout(
	name string, h core.InfoHash, peer *core.PeerInfo) ([]*core.PeerInfo, error) {

	if peer.Complete {
		// If the peer is announcing as complete, don't return a peer handout since
		// the peer does not need it.
		return nil, nil
	}
	var errs []error
	peers, err := s.peerStore.GetPeers(h, s.config.PeerHandoutLimit)
	if err != nil {
		errs = append(errs, fmt.Errorf("peer storage error: %s", err))
	}
	origins, err := s.peerStore.GetOrigins(h)
	if err != nil {
		tryUpdate := true
		if err != storage.ErrNoOrigins {
			errs = append(errs, fmt.Errorf("origin peer storage error: %s", err))
			tryUpdate = false
		}
		origins, err = s.fetchOrigins(name)
		if err != nil {
			errs = append(errs, fmt.Errorf("origin lookup error: %s", err))
			tryUpdate = false
		}
		if tryUpdate {
			if err := s.peerStore.UpdateOrigins(h, origins); err != nil {
				log.With("hash", h).Errorf("Error upserting origins: %s", err)
			}
		}
	}
	err = errutil.Join(errs)
	if err != nil {
		log.Errorf("Error getting peers: %s", err)
	}
	for _, origin := range origins {
		peers = append(peers, origin)
	}
	if len(peers) == 0 {
		if err != nil {
			return nil, handler.Errorf("error getting peers: %s", err)
		}
		return nil, handler.ErrorStatus(http.StatusNotFound)
	}

	err = s.policy.AssignPeerPriority(peer, peers)
	if err != nil {
		return nil, handler.Errorf("assign peer priority: %s", err)
	}

	// TODO(codyg): Just make this sort peers since storage already samples via limit.
	peers, err = s.policy.SamplePeers(peers, len(peers))
	if err != nil {
		return nil, handler.Errorf("sample peers: %s", err)
	}

	return peers, nil
}

func (s *Server) fetchOrigins(name string) ([]*core.PeerInfo, error) {
	d, err := core.NewSHA256DigestFromHex(name)
	if err != nil {
		return nil, fmt.Errorf("new digest: %s", err)
	}
	octxs, err := s.originCluster.Owners(d)
	if err != nil {
		return nil, err
	}
	var origins []*core.PeerInfo
	for _, octx := range octxs {
		origins = append(origins, core.PeerInfoFromContext(octx, true))
	}
	return origins, nil
}
<commit_msg>Record tracker hand-out stats<commit_after>package trackerserver

import (
	"encoding/json"
	"fmt"
	"net/http"

	"code.uber.internal/infra/kraken/core"
	"code.uber.internal/infra/kraken/tracker/announceclient"
	"code.uber.internal/infra/kraken/tracker/storage"
	"code.uber.internal/infra/kraken/utils/errutil"
	"code.uber.internal/infra/kraken/utils/handler"
	"code.uber.internal/infra/kraken/utils/log"
)

const (
	_seederHandoutGauge    = "seeders_handed_out"
	_seederHandoutPctGauge = "seeders_handed_out_pct"
)

func (s *Server) announceHandler(w http.ResponseWriter, r *http.Request) error {
	req := new(announceclient.Request)
	if err := json.NewDecoder(r.Body).Decode(req); err != nil {
		return handler.Errorf("json decode request: %s", err)
	}
	if err := s.peerStore.UpdatePeer(req.InfoHash, req.Peer); err != nil {
		log.With(
			"hash", req.InfoHash,
			"peer_id", req.Peer.PeerID).Errorf("Error updating peer: %s", err)
	}
	peers, err := s.getPeerHandout(req.Name, req.InfoHash, req.Peer)
	if err != nil {
		return err
	}
	resp := &announceclient.Response{
		Peers:    peers,
		Interval: s.config.AnnounceInterval,
	}
	if err := json.NewEncoder(w).Encode(resp); err != nil {
		return handler.Errorf("json encode response: %s", err)
	}
	return nil
}

func (s *Server) getPeerHandout(
	name string, h core.InfoHash, peer *core.PeerInfo) ([]*core.PeerInfo, error) {

	if peer.Complete {
		// If the peer is announcing as complete, don't return a peer handout since
		// the peer does not need it.
		return nil, nil
	}
	var errs []error
	peers, err := s.peerStore.GetPeers(h, s.config.PeerHandoutLimit)
	if err != nil {
		errs = append(errs, fmt.Errorf("peer storage error: %s", err))
	}
	origins, err := s.peerStore.GetOrigins(h)
	if err != nil {
		tryUpdate := true
		if err != storage.ErrNoOrigins {
			errs = append(errs, fmt.Errorf("origin peer storage error: %s", err))
			tryUpdate = false
		}
		origins, err = s.fetchOrigins(name)
		if err != nil {
			errs = append(errs, fmt.Errorf("origin lookup error: %s", err))
			tryUpdate = false
		}
		if tryUpdate {
			if err := s.peerStore.UpdateOrigins(h, origins); err != nil {
				log.With("hash", h).Errorf("Error upserting origins: %s", err)
			}
		}
	}
	err = errutil.Join(errs)
	if err != nil {
		log.Errorf("Error getting peers: %s", err)
	}
	for _, origin := range origins {
		peers = append(peers, origin)
	}
	if len(peers) == 0 {
		if err != nil {
			return nil, handler.Errorf("error getting peers: %s", err)
		}
		return nil, handler.ErrorStatus(http.StatusNotFound)
	}

	err = s.policy.AssignPeerPriority(peer, peers)
	if err != nil {
		return nil, handler.Errorf("assign peer priority: %s", err)
	}

	// TODO(codyg): Just make this sort peers since storage already samples via limit.
	peers, err = s.policy.SamplePeers(peers, len(peers))
	if err != nil {
		return nil, handler.Errorf("sample peers: %s", err)
	}

	var peerCount int
	var seederCount int
	for _, peer := range peers {
		if !peer.Origin {
			peerCount++
			if peer.Complete {
				seederCount++
			}
		}
	}

	s.stats.Gauge(_seederHandoutGauge).Update(float64(seederCount))
	s.stats.Gauge(_seederHandoutPctGauge).Update(float64(seederCount) / float64(peerCount))

	return peers, nil
}

func (s *Server) fetchOrigins(name string) ([]*core.PeerInfo, error) {
	d, err := core.NewSHA256DigestFromHex(name)
	if err != nil {
		return nil, fmt.Errorf("new digest: %s", err)
	}
	octxs, err := s.originCluster.Owners(d)
	if err != nil {
		return nil, err
	}
	var origins []*core.PeerInfo
	for _, octx := range octxs {
		origins = append(origins, core.PeerInfoFromContext(octx, true))
	}
	return origins, nil
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/didip/tollbooth"
	"time"
	"./middleware"
	"github.com/gin-gonic/gin"
)

//func HelloHandler(w http.ResponseWriter, req *http.Request) {
//	w.Write([]byte("Hello, World!"))
//}
//
//func main() {
//	// Create a request limiter per handler.
//	http.Handle("/", tollbooth.LimitFuncHandler(tollbooth.NewLimiter(1, 1* time.Second), HelloHandler))
//	http.ListenAndServe(":12345", nil)
//}

func main() {
	r := gin.New()

	// Create a limiter struct.
	limiter := tollbooth.NewLimiter(1, time.Second)
	r.GET("/", middleware.LimitHandler(limiter), func(c *gin.Context) {
		c.String(200, "Hello, world!")
	})

	r.Run(":12345")
}
<commit_msg>Update main.go<commit_after>package main

import (
	"github.com/didip/tollbooth"
	"time"
	"./middleware"
	"github.com/gin-gonic/gin"
)

//func HelloHandler(w http.ResponseWriter, req *http.Request) {
//	w.Write([]byte("Hello, World!"))
//}
//
//func main() {
//	// Create a request limiter per handler.
//	http.Handle("/", tollbooth.LimitFuncHandler(tollbooth.NewLimiter(1, 1* time.Second), HelloHandler))
//	http.ListenAndServe(":12345", nil)
//}

func main() {
	r := gin.New()

	// Create a limiter struct.
	limiter := tollbooth.NewLimiter(1, time.Second)
	r.GET("/", middleware.LimitHandler(limiter), func(c *gin.Context) {
		c.String(200, "Hello, world!")
	})

	r.Run(":12345")
}

// 其他参考 https://github.com/VojtechVitek/ratelimit
<|endoftext|>


<commit_before>package gitosis

import (
	"github.com/timeredbull/tsuru/config"
	"github.com/timeredbull/tsuru/log"
	"io/ioutil"
	. "launchpad.net/gocheck"
	stdlog "log"
	"os"
	"os/exec"
	"path"
	"strings"
	"testing"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
	gitRoot     string
	gitosisBare string
	gitosisRepo string
	logFile     *os.File
}

var _ = Suite(&S{})

type isInGitosisChecker struct{}

func (c *isInGitosisChecker) Info() *CheckerInfo {
	return &CheckerInfo{Name: "IsInGitosis", Params: []string{"str"}}
}

func (c *isInGitosisChecker) Check(params []interface{}, names []string) (bool, string) {
	if len(params) != 1 {
		return false, "you should provide one string parameter"
	}
	str, ok := params[0].(string)
	if !ok {
		return false, "the parameter should be a string"
	}
	gitosisRepo, err := config.GetString("git:gitosis-repo")
	if err != nil {
		return false, "failed to get config"
	}
	path := path.Join(gitosisRepo, "gitosis.conf")
	f, err := os.Open(path)
	if err != nil {
		return false, err.Error()
	}
	defer f.Close()
	content, err := ioutil.ReadAll(f)
	if err != nil {
		return false, err.Error()
	}
	return strings.Contains(string(content), str), ""
}

var IsInGitosis, NotInGitosis Checker = &isInGitosisChecker{}, Not(IsInGitosis)

func (s *S) SetUpSuite(c *C) {
	err := config.ReadConfigFile("../etc/tsuru.conf")
	c.Assert(err, IsNil)
	s.gitRoot, err = config.GetString("git:root")
	c.Assert(err, IsNil)
	s.gitosisBare, err = config.GetString("git:gitosis-bare")
	c.Assert(err, IsNil)
	s.gitosisRepo, err = config.GetString("git:gitosis-repo")
	err = os.RemoveAll(s.gitRoot)
	c.Assert(err, IsNil)
	err = os.MkdirAll(s.gitRoot, 0777)
	c.Assert(err, IsNil)
	err = exec.Command("git", "init", "--bare", s.gitosisBare).Run()
	c.Assert(err, IsNil)
	err = exec.Command("git", "clone", s.gitosisBare, s.gitosisRepo).Run()
	c.Assert(err, IsNil)
	s.logFile, err = os.Create("/tmp/tsuru-tests.log")
	c.Assert(err, IsNil)
	log.Target = stdlog.New(s.logFile, "[tsuru-tests]", stdlog.LstdFlags|stdlog.Llongfile)
}

func (s *S) SetUpTest(c *C) {
	fpath := path.Join(s.gitosisRepo, "gitosis.conf")
	f, err := os.Create(fpath)
	c.Assert(err, IsNil)
	f.Close()
}

func (s *S) TearDownSuite(c *C) {
	defer s.logFile.Close()
	err := os.RemoveAll(s.gitRoot)
	c.Assert(err, IsNil)
}

func (s *S) TearDownTest(c *C) {
	_, err := runGit("rm", "gitosis.conf")
	if err == nil {
		err = pushToGitosis("removing test file")
		c.Assert(err, IsNil)
	}
}

func (s *S) lastBareCommit(c *C) string {
	bareOutput, err := exec.Command("git", "--git-dir="+s.gitosisBare, "log", "-1", "--pretty=format:%s").CombinedOutput()
	c.Assert(err, IsNil)
	return string(bareOutput)
}
<commit_msg>gitosis: improving test fixtures<commit_after>package gitosis

import (
	"github.com/timeredbull/tsuru/config"
	"github.com/timeredbull/tsuru/log"
	"io/ioutil"
	. "launchpad.net/gocheck"
	stdlog "log"
	"os"
	"os/exec"
	"path"
	"strings"
	"testing"
)

func Test(t *testing.T) { TestingT(t) }

type S struct {
	gitRoot     string
	gitosisBare string
	gitosisRepo string
	logFile     *os.File
}

var _ = Suite(&S{})

type isInGitosisChecker struct{}

func (c *isInGitosisChecker) Info() *CheckerInfo {
	return &CheckerInfo{Name: "IsInGitosis", Params: []string{"str"}}
}

func (c *isInGitosisChecker) Check(params []interface{}, names []string) (bool, string) {
	if len(params) != 1 {
		return false, "you should provide one string parameter"
	}
	str, ok := params[0].(string)
	if !ok {
		return false, "the parameter should be a string"
	}
	gitosisRepo, err := config.GetString("git:gitosis-repo")
	if err != nil {
		return false, "failed to get config"
	}
	path := path.Join(gitosisRepo, "gitosis.conf")
	f, err := os.Open(path)
	if err != nil {
		return false, err.Error()
	}
	defer f.Close()
	content, err := ioutil.ReadAll(f)
	if err != nil {
		return false, err.Error()
	}
	return strings.Contains(string(content), str), ""
}

var IsInGitosis, NotInGitosis Checker = &isInGitosisChecker{}, Not(IsInGitosis)

func (s *S) SetUpSuite(c *C) {
	err := config.ReadConfigFile("../etc/tsuru.conf")
	c.Assert(err, IsNil)
	s.gitRoot, err = config.GetString("git:root")
	c.Assert(err, IsNil)
	s.gitosisBare, err = config.GetString("git:gitosis-bare")
	c.Assert(err, IsNil)
	s.gitosisRepo, err = config.GetString("git:gitosis-repo")
	err = os.RemoveAll(s.gitRoot)
	c.Assert(err, IsNil)
	err = os.MkdirAll(s.gitRoot, 0777)
	c.Assert(err, IsNil)
	err = exec.Command("git", "init", "--bare", s.gitosisBare).Run()
	c.Assert(err, IsNil)
	err = exec.Command("git", "clone", s.gitosisBare, s.gitosisRepo).Run()
	c.Assert(err, IsNil)
	s.logFile, err = os.Create("/tmp/tsuru-tests.log")
	c.Assert(err, IsNil)
	log.Target = stdlog.New(s.logFile, "[tsuru-tests]", stdlog.LstdFlags|stdlog.Llongfile)
}

func (s *S) SetUpTest(c *C) {
	fpath := path.Join(s.gitosisRepo, "gitosis.conf")
	f, err := os.Create(fpath)
	c.Assert(err, IsNil)
	f.Close()
	pushToGitosis("added gitosis test file")
}

func (s *S) TearDownSuite(c *C) {
	defer s.logFile.Close()
	err := os.RemoveAll(s.gitRoot)
	c.Assert(err, IsNil)
}

func (s *S) TearDownTest(c *C) {
	_, err := runGit("rm", "gitosis.conf")
	c.Assert(err, IsNil)
	err = pushToGitosis("removing test file")
	c.Assert(err, IsNil)
}

func (s *S) lastBareCommit(c *C) string {
	bareOutput, err := exec.Command("git", "--git-dir="+s.gitosisBare, "log", "-1", "--pretty=format:%s").CombinedOutput()
	c.Assert(err, IsNil)
	return string(bareOutput)
}
<|endoftext|>


<commit_before>package engine

import (
	"fmt"

	"github.com/keybase/client/go/libkb"
)

// Identify is an engine to identify a user.
type Identify struct {
	arg       *IdentifyArg
	user      *libkb.User
	userExpr  libkb.AssertionExpression
	outcome   *libkb.IdentifyOutcome
	trackInst *libkb.TrackInstructions
}

type IdentifyArg struct {
	TargetUsername string // The user being identified, leave blank to identify self
	WithTracking   bool   // true if want tracking statement for logged in user on TargetUsername

	// When tracking is being performed, the identify engine is used with a tracking ui.
	// These options are sent to the ui based on command line options.
	// For normal identify, safe to leave these in their default zero state.
	TrackOptions TrackOptions
}

func NewIdentifyArg(targetUsername string, withTracking bool) *IdentifyArg {
	return &IdentifyArg{
		TargetUsername: targetUsername,
		WithTracking:   withTracking,
	}
}

func NewIdentifyTrackArg(targetUsername string, withTracking bool, options TrackOptions) *IdentifyArg {
	return &IdentifyArg{
		TargetUsername: targetUsername,
		WithTracking:   withTracking,
		TrackOptions:   options,
	}
}

// NewIdentify creates a Identify engine.
func NewIdentify(arg *IdentifyArg) *Identify {
	return &Identify{arg: arg}
}

// Name is the unique engine name.
func (e *Identify) Name() string {
	return "Identify"
}

// GetPrereqs returns the engine prereqs.
func (e *Identify) GetPrereqs() EnginePrereqs {
	// if WithTracking is on, we need to be logged in
	return EnginePrereqs{Session: e.arg.WithTracking}
}

// RequiredUIs returns the required UIs.
func (e *Identify) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.IdentifyUIKind,
	}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *Identify) SubConsumers() []libkb.UIConsumer {
	return nil
}

// Run starts the engine.
func (e *Identify) Run(ctx *Context) error {
	ok, err := IsLoggedIn()
	if err != nil {
		return err
	}
	if ok {
		// logged in, so turn on WithTracking as there's no reason not to:
		e.arg.WithTracking = true
	}

	if err := e.loadUser(); err != nil {
		return err
	}

	ctx.IdentifyUI.Start(e.user.GetName())

	e.outcome, err = e.run(ctx)
	if err != nil {
		return err
	}

	// set the flags in the outcome with the track options to
	// inform the ui what to do with the remote tracking prompt:
	e.outcome.LocalOnly = e.arg.TrackOptions.TrackLocalOnly
	e.outcome.ApproveRemote = e.arg.TrackOptions.TrackApprove

	tmp, err := ctx.IdentifyUI.FinishAndPrompt(e.outcome.Export())
	if err != nil {
		return err
	}
	fpr := libkb.ImportFinishAndPromptRes(tmp)
	e.trackInst = &fpr

	return nil
}

func (e *Identify) User() *libkb.User {
	return e.user
}

func (e *Identify) Outcome() *libkb.IdentifyOutcome {
	return e.outcome
}

func (e *Identify) TrackInstructions() *libkb.TrackInstructions {
	return e.trackInst
}

func (e *Identify) run(ctx *Context) (*libkb.IdentifyOutcome, error) {
	res := libkb.NewIdentifyOutcome(e.arg.WithTracking)
	is := libkb.NewIdentifyState(res, e.user)

	if e.arg.WithTracking {
		me, err := libkb.LoadMe(libkb.LoadUserArg{})
		if err != nil {
			return nil, err
		}

		if e.user.Equal(*me) {
			return nil, libkb.SelfTrackError{}
		}

		tlink, err := me.GetTrackingStatementFor(e.user.GetName(), e.user.GetUid())
		if err != nil {
			return nil, err
		}
		if tlink != nil {
			is.Track = libkb.NewTrackLookup(tlink)
			res.TrackUsed = is.Track
		}
	}

	ctx.IdentifyUI.ReportLastTrack(libkb.ExportTrackSummary(is.Track))

	G.Log.Debug("+ Identify(%s)", e.user.GetName())

	for _, bundle := range e.user.GetActivePgpKeys(true) {
		fokid := libkb.GenericKeyToFOKID(bundle)
		var diff libkb.TrackDiff
		if is.Track != nil {
			diff = is.Track.ComputeKeyDiff(&fokid)
			// XXX this is probably a bug now that there are multiple pgp keys
			res.KeyDiff = diff
		}
		ctx.IdentifyUI.DisplayKey(fokid.Export(), libkb.ExportTrackDiff(diff))
	}

	is.InitResultList()
	is.ComputeTrackDiffs()
	is.ComputeDeletedProofs()

	ctx.IdentifyUI.LaunchNetworkChecks(res.ExportToUncheckedIdentity(), e.user.Export())
	e.user.IdTable.Identify(is, ctx.IdentifyUI)

	if !e.userExpr.MatchSet(*e.user.ToOkProofSet()) {
		return nil, fmt.Errorf("User %s didn't match given assertion", e.user.GetName())
	}

	G.Log.Debug("- Identify(%s)", e.user.GetName())

	return res, nil
}

func (e *Identify) loadUser() error {
	arg, err := e.loadUserArg()
	if err != nil {
		return err
	}

	u, err := libkb.LoadUser(*arg)
	if err != nil {
		return err
	}
	e.user = u

	return nil
}

func (e *Identify) loadUserArg() (*libkb.LoadUserArg, error) {
	if len(e.arg.TargetUsername) == 0 {
		// loading self
		return &libkb.LoadUserArg{Self: true}, nil
	}

	// Use assertions for everything:

	// Parse assertion but don't allow OR operators, only
	// AND operators
	expr, err := libkb.AssertionParseAndOnly(e.arg.TargetUsername)
	if err != nil {
		return nil, err
	}
	e.userExpr = expr

	// Next, pop off the 'best' assertion and load the user by it.
	// That is, it might be the keybase assertion (if there), or otherwise,
	// something that's unique like Twitter or Github, and lastly,
	// something like DNS that is more likely ambiguous...
	b := e.findBestComponent(expr)
	if len(b) == 0 {
		return nil, fmt.Errorf("Cannot lookup user with %q", e.arg.TargetUsername)
	}

	return &libkb.LoadUserArg{Name: b}, nil
}

func (e *Identify) findBestComponent(expr libkb.AssertionExpression) string {
	urls := make([]libkb.AssertionUrl, 0, 1)
	urls = expr.CollectUrls(urls)
	if len(urls) == 0 {
		return ""
	}

	var uid, kb, soc, fp libkb.AssertionUrl

	for _, u := range urls {
		if u.IsUid() {
			uid = u
			break
		} else if u.IsKeybase() {
			kb = u
		} else if u.IsFingerprint() && fp == nil {
			fp = u
		} else if u.IsSocial() && soc == nil {
			soc = u
		}
	}

	order := []libkb.AssertionUrl{uid, kb, fp, soc, urls[0]}
	for _, p := range order {
		if p != nil {
			return p.String()
		}
	}
	return ""
}
<commit_msg>fixed self id<commit_after>package engine

import (
	"fmt"

	"github.com/keybase/client/go/libkb"
)

// Identify is an engine to identify a user.
type Identify struct {
	arg       *IdentifyArg
	user      *libkb.User
	userExpr  libkb.AssertionExpression
	outcome   *libkb.IdentifyOutcome
	trackInst *libkb.TrackInstructions
}

type IdentifyArg struct {
	TargetUsername string // The user being identified, leave blank to identify self
	WithTracking   bool   // true if want tracking statement for logged in user on TargetUsername

	// When tracking is being performed, the identify engine is used with a tracking ui.
	// These options are sent to the ui based on command line options.
	// For normal identify, safe to leave these in their default zero state.
	TrackOptions TrackOptions
}

func NewIdentifyArg(targetUsername string, withTracking bool) *IdentifyArg {
	return &IdentifyArg{
		TargetUsername: targetUsername,
		WithTracking:   withTracking,
	}
}

func NewIdentifyTrackArg(targetUsername string, withTracking bool, options TrackOptions) *IdentifyArg {
	return &IdentifyArg{
		TargetUsername: targetUsername,
		WithTracking:   withTracking,
		TrackOptions:   options,
	}
}

func (ia *IdentifyArg) SelfID() bool {
	return len(ia.TargetUsername) == 0
}

// NewIdentify creates a Identify engine.
func NewIdentify(arg *IdentifyArg) *Identify {
	return &Identify{arg: arg}
}

// Name is the unique engine name.
func (e *Identify) Name() string {
	return "Identify"
}

// GetPrereqs returns the engine prereqs.
func (e *Identify) GetPrereqs() EnginePrereqs {
	// if WithTracking is on, we need to be logged in
	return EnginePrereqs{Session: e.arg.WithTracking}
}

// RequiredUIs returns the required UIs.
func (e *Identify) RequiredUIs() []libkb.UIKind {
	return []libkb.UIKind{
		libkb.IdentifyUIKind,
	}
}

// SubConsumers returns the other UI consumers for this engine.
func (e *Identify) SubConsumers() []libkb.UIConsumer {
	return nil
}

// Run starts the engine.
func (e *Identify) Run(ctx *Context) error {
	ok, err := IsLoggedIn()
	if err != nil {
		return err
	}
	if ok && !e.arg.SelfID() {
		// logged in, so turn on WithTracking as there's no reason not to,
		// unless doing a self id:
		e.arg.WithTracking = true
	}

	if err := e.loadUser(); err != nil {
		return err
	}

	ctx.IdentifyUI.Start(e.user.GetName())

	e.outcome, err = e.run(ctx)
	if err != nil {
		return err
	}

	// set the flags in the outcome with the track options to
	// inform the ui what to do with the remote tracking prompt:
	e.outcome.LocalOnly = e.arg.TrackOptions.TrackLocalOnly
	e.outcome.ApproveRemote = e.arg.TrackOptions.TrackApprove

	tmp, err := ctx.IdentifyUI.FinishAndPrompt(e.outcome.Export())
	if err != nil {
		return err
	}
	fpr := libkb.ImportFinishAndPromptRes(tmp)
	e.trackInst = &fpr

	return nil
}

func (e *Identify) User() *libkb.User {
	return e.user
}

func (e *Identify) Outcome() *libkb.IdentifyOutcome {
	return e.outcome
}

func (e *Identify) TrackInstructions() *libkb.TrackInstructions {
	return e.trackInst
}

func (e *Identify) run(ctx *Context) (*libkb.IdentifyOutcome, error) {
	res := libkb.NewIdentifyOutcome(e.arg.WithTracking)
	is := libkb.NewIdentifyState(res, e.user)

	if e.arg.WithTracking {
		me, err := libkb.LoadMe(libkb.LoadUserArg{})
		if err != nil {
			return nil, err
		}

		if e.user.Equal(*me) {
			return nil, libkb.SelfTrackError{}
		}

		tlink, err := me.GetTrackingStatementFor(e.user.GetName(), e.user.GetUid())
		if err != nil {
			return nil, err
		}
		if tlink != nil {
			is.Track = libkb.NewTrackLookup(tlink)
			res.TrackUsed = is.Track
		}
	}

	ctx.IdentifyUI.ReportLastTrack(libkb.ExportTrackSummary(is.Track))

	G.Log.Debug("+ Identify(%s)", e.user.GetName())

	for _, bundle := range e.user.GetActivePgpKeys(true) {
		fokid := libkb.GenericKeyToFOKID(bundle)
		var diff libkb.TrackDiff
		if is.Track != nil {
			diff = is.Track.ComputeKeyDiff(&fokid)
			// XXX this is probably a bug now that there are multiple pgp keys
			res.KeyDiff = diff
		}
		ctx.IdentifyUI.DisplayKey(fokid.Export(), libkb.ExportTrackDiff(diff))
	}

	is.InitResultList()
	is.ComputeTrackDiffs()
	is.ComputeDeletedProofs()

	ctx.IdentifyUI.LaunchNetworkChecks(res.ExportToUncheckedIdentity(), e.user.Export())
	e.user.IdTable.Identify(is, ctx.IdentifyUI)

	if !e.userExpr.MatchSet(*e.user.ToOkProofSet()) {
		return nil, fmt.Errorf("User %s didn't match given assertion", e.user.GetName())
	}

	G.Log.Debug("- Identify(%s)", e.user.GetName())

	return res, nil
}

func (e *Identify) loadUser() error {
	arg, err := e.loadUserArg()
	if err != nil {
		return err
	}

	u, err := libkb.LoadUser(*arg)
	if err != nil {
		return err
	}
	e.user = u

	if arg.Self {
		// if this was a self load, need to load an assertion expression
		// now that we have the username
		if err := e.loadExpr(e.user.GetName()); err != nil {
			return err
		}
	}

	return nil
}

func (e *Identify) loadUserArg() (*libkb.LoadUserArg, error) {
	if e.arg.SelfID() {
		// loading self
		return &libkb.LoadUserArg{Self: true}, nil
	}

	// Use assertions for everything:
	if err := e.loadExpr(e.arg.TargetUsername); err != nil {
		return nil, err
	}

	// Next, pop off the 'best' assertion and load the user by it.
	// That is, it might be the keybase assertion (if there), or otherwise,
	// something that's unique like Twitter or Github, and lastly,
	// something like DNS that is more likely ambiguous...
	b := e.findBestComponent(e.userExpr)
	if len(b) == 0 {
		return nil, fmt.Errorf("Cannot lookup user with %q", e.arg.TargetUsername)
	}

	return &libkb.LoadUserArg{Name: b}, nil
}

func (e *Identify) loadExpr(assertion string) error {
	// Parse assertion but don't allow OR operators, only AND operators
	expr, err := libkb.AssertionParseAndOnly(assertion)
	if err != nil {
		return fmt.Errorf("assertion parse error: %s", err)
	}
	e.userExpr = expr
	return nil
}

func (e *Identify) findBestComponent(expr libkb.AssertionExpression) string {
	urls := make([]libkb.AssertionUrl, 0, 1)
	urls = expr.CollectUrls(urls)
	if len(urls) == 0 {
		return ""
	}

	var uid, kb, soc, fp libkb.AssertionUrl

	for _, u := range urls {
		if u.IsUid() {
			uid = u
			break
		} else if u.IsKeybase() {
			kb = u
		} else if u.IsFingerprint() && fp == nil {
			fp = u
		} else if u.IsSocial() && soc == nil {
			soc = u
		}
	}

	order := []libkb.AssertionUrl{uid, kb, fp, soc, urls[0]}
	for _, p := range order {
		if p != nil {
			return p.String()
		}
	}
	return ""
}
<|endoftext|>


<commit_before>package graph

import (
	"github.com/twmb/algoimpl/go/graph/lite"
	"math/rand"
	"sync"
)

const (
	dequeued = ^(1<<31 - 1)
	unseen   = 0
	seen     = 1
)

// O(V + E). It does not matter to traverse back
// on a bidirectional edge, because any vertex dfs is
// recursing on is marked as visited and won't be visited
// again anyway.
func (g *Graph) dfs(node *node, finishList *[]Node) {
	node.state = seen
	for _, edge := range node.edges {
		if edge.end.state == unseen {
			edge.end.parent = node
			g.dfs(edge.end, finishList)
		}
	}
	*finishList = append(*finishList, node.container)
}

func (g *Graph) dfsReversedEdges(node *node, finishList *[]Node) {
	node.state = seen
	for _, edge := range node.reversedEdges {
		if edge.end.state == unseen {
			edge.end.parent = node
			g.dfsReversedEdges(edge.end, finishList)
		}
	}
	*finishList = append(*finishList, node.container)
}

// TopologicalSort topoligically sorts a directed acyclic graph.
// If the graph is cyclic, the sort order will change
// based on which node the sort starts on.
//
// The StronglyConnectedComponents function can be used to determine if a graph has cycles.
func (g *Graph) TopologicalSort() []Node {
	sorted := make([]Node, 0, len(g.nodes))
	// sort preorder (first jacket, then shirt)
	for _, node := range g.nodes {
		if node.state == unseen {
			g.dfs(node, &sorted)
		}
	}
	// now make post order for correct sort (jacket follows shirt). O(V)
	length := len(sorted)
	for i := 0; i < length/2; i++ {
		sorted[i], sorted[length-i-1] = sorted[length-i-1], sorted[i]
	}
	return sorted
}

// Reverse returns reversed copy of the directed graph g.
// This function can be used to copy an undirected graph.
func (g *Graph) Reverse() *Graph {
	reversed := New(Directed)
	if g.Kind == Undirected {
		reversed = New(Undirected)
	}
	// O(V)
	for _ = range g.nodes {
		reversed.MakeNode()
	}
	// O(V + E)
	for _, node := range g.nodes {
		for _, edge := range node.edges {
			reversed.MakeEdge(reversed.nodes[edge.end.index].container,
				reversed.nodes[node.index].container)
		}
	}
	return reversed
}

// StronglyConnectedComponents returns a slice of strongly connected nodes in a directed graph.
// If used on an undirected graph, this function returns nil.
func (g *Graph) StronglyConnectedComponents() [][]Node {
	if g.Kind == Undirected {
		return nil
	}
	components := make([][]Node, 0)
	finishOrder := g.TopologicalSort()
	for i := range finishOrder {
		finishOrder[i].node.state = unseen
	}
	for _, sink := range finishOrder {
		if g.nodes[sink.node.index].state == unseen {
			component := make([]Node, 0)
			g.dfsReversedEdges(g.nodes[sink.node.index], &component)
			components = append(components, component)
		}
	}
	return components
}

// RandMinimumCut runs Kargers algorithm to find a random minimum cut
// on the graph. If iterations is < 1, this will return an empty slice.
// Otherwise, it returns a slice of the edges crossing the best minimum
// cut found in any iteration. Call rand.Seed() before using this function.
//
// This function takes a number of iterations to start concurrently. If
// concurrent is <= 1, it will run one iteration at a time.
//
// If the graph is Directed, this will return a cut of edges in both directions.
// If the graph is Undirected, this will return a proper min cut.
func (g *Graph) RandMinimumCut(iterations, concurrent int) []Edge {
	var mutex sync.Mutex
	doneChan := make(chan struct{}, iterations)
	if concurrent < 1 {
		concurrent = 1
	}
	sem := make(chan struct{}, concurrent)
	for i := 0; i < concurrent; i++ {
		sem <- struct{}{}
	}
	// make a lite slice of the edges
	var baseAllEdges []lite.Edge
	for n := range g.nodes {
		for _, edge := range g.nodes[n].edges {
			if g.Kind == Undirected && n < edge.end.index {
				continue
			}
			baseAllEdges = append(baseAllEdges, lite.Edge{Start: n,
				End: edge.end.index, S: g.nodes[n], E: edge})
		}
	}

	minCutLite := make([]lite.Edge, len(baseAllEdges))

	for iter := 0; iter < iterations; iter++ {
		<-sem
		go func() {
			nodecount := len(g.nodes)
			allEdges := make([]lite.Edge, len(baseAllEdges))
			copy(allEdges, baseAllEdges)
			// shuffle for random edge removal order
			shuffle(allEdges)
			for nodecount > 2 {
				// remove first edge, keep the start node, collapse the end node
				// anything that points to the collapsing node now points to the keep node
				// anything that starts at the collapsing node now starts at the keep node
				keep := allEdges[len(allEdges)-1].Start
				remove := allEdges[len(allEdges)-1].End
				allEdges = allEdges[:len(allEdges)-1]
				for e := 0; e < len(allEdges); e++ {
					if allEdges[e].Start == remove {
						allEdges[e].Start = keep
					}
					if allEdges[e].End == remove {
						allEdges[e].End = keep
					}
					// remove the edge if it self looped
					if allEdges[e].Start == allEdges[e].End {
						allEdges[e] = allEdges[len(allEdges)-1]
						allEdges = allEdges[:len(allEdges)-1]
						e--
					}
				}
				// every edge removed removes a node
				nodecount--
			}

			mutex.Lock()
			if iter == 0 || len(allEdges) < len(minCutLite) {
				minCutLite = make([]lite.Edge, len(allEdges))
				copy(minCutLite, allEdges)
			}
			mutex.Unlock()

			doneChan <- struct{}{}
			sem <- struct{}{}
		}()
	}
	for iter := 0; iter < iterations; iter++ {
		<-doneChan
	}

	minCut := make([]Edge, len(minCutLite))
	for i := range minCutLite {
		start := minCutLite[i].S.(*node)
		edge := minCutLite[i].E.(edge)
		minCut[i] = Edge{Weight: edge.weight, Start: start.container,
			End: edge.end.container}
	}
	return minCut
}

// Fischer-Yates shuffle
func shuffle(edges []lite.Edge) {
	for i := len(edges) - 1; i > 0; i-- {
		j := rand.Intn(i + 1)
		edges[j], edges[i] = edges[i], edges[j]
	}
}

// MinimumSpanningTree will return the edges corresponding to the
// minimum spanning tree in the graph based off of edge weight values.
// This will return nil for a directed graph.
func (g *Graph) MinimumSpanningTree() []Edge {
	if g.Kind == Directed {
		return nil
	}
	// create priority queue for vertices
	// node.data is the index into the heap
	// node.state is the weight of an edge to its parent
	nodesBase := nodeSlice(make([]*node, len(g.nodes)))
	copy(nodesBase, g.nodes)
	for i := range nodesBase {
		nodesBase[i].state = 1<<31 - 1
		nodesBase[i].data = i
	}
	nodesBase[0].state = 0
	nodesBase[0].parent = nil
	nodes := &nodesBase
	nodes.heapInit()

	for len(*nodes) > 0 {
		min := nodes.pop()
		for _, edge := range min.edges {
			v := edge.end // get the other side of the edge
			if nodes.heapContains(v) && edge.weight < v.state {
				v.parent = min
				nodes.update(v.data, edge.weight)
			}
		}
	}

	mst := make([]Edge, 0, len(g.nodes))
	for i := range g.nodes {
		if g.nodes[i].parent != nil {
			mst = append(mst, Edge{Weight: g.nodes[i].state,
				Start: g.nodes[i].container, End: g.nodes[i].parent.container})
		}
	}

	return mst
}
<commit_msg>mst has (# of nodes -1) edges, not (# of nodes) - minor fix<commit_after>package graph

import (
	"github.com/twmb/algoimpl/go/graph/lite"
	"math/rand"
	"sync"
)

const (
	dequeued = ^(1<<31 - 1)
	unseen   = 0
	seen     = 1
)

// O(V + E). It does not matter to traverse back
// on a bidirectional edge, because any vertex dfs is
// recursing on is marked as visited and won't be visited
// again anyway.
func (g *Graph) dfs(node *node, finishList *[]Node) {
	node.state = seen
	for _, edge := range node.edges {
		if edge.end.state == unseen {
			edge.end.parent = node
			g.dfs(edge.end, finishList)
		}
	}
	*finishList = append(*finishList, node.container)
}

func (g *Graph) dfsReversedEdges(node *node, finishList *[]Node) {
	node.state = seen
	for _, edge := range node.reversedEdges {
		if edge.end.state == unseen {
			edge.end.parent = node
			g.dfsReversedEdges(edge.end, finishList)
		}
	}
	*finishList = append(*finishList, node.container)
}

// TopologicalSort topoligically sorts a directed acyclic graph.
// If the graph is cyclic, the sort order will change
// based on which node the sort starts on.
//
// The StronglyConnectedComponents function can be used to determine if a graph has cycles.
func (g *Graph) TopologicalSort() []Node {
	sorted := make([]Node, 0, len(g.nodes))
	// sort preorder (first jacket, then shirt)
	for _, node := range g.nodes {
		if node.state == unseen {
			g.dfs(node, &sorted)
		}
	}
	// now make post order for correct sort (jacket follows shirt). O(V)
	length := len(sorted)
	for i := 0; i < length/2; i++ {
		sorted[i], sorted[length-i-1] = sorted[length-i-1], sorted[i]
	}
	return sorted
}

// Reverse returns reversed copy of the directed graph g.
// This function can be used to copy an undirected graph.
func (g *Graph) Reverse() *Graph {
	reversed := New(Directed)
	if g.Kind == Undirected {
		reversed = New(Undirected)
	}
	// O(V)
	for _ = range g.nodes {
		reversed.MakeNode()
	}
	// O(V + E)
	for _, node := range g.nodes {
		for _, edge := range node.edges {
			reversed.MakeEdge(reversed.nodes[edge.end.index].container,
				reversed.nodes[node.index].container)
		}
	}
	return reversed
}

// StronglyConnectedComponents returns a slice of strongly connected nodes in a directed graph.
// If used on an undirected graph, this function returns nil.
func (g *Graph) StronglyConnectedComponents() [][]Node {
	if g.Kind == Undirected {
		return nil
	}
	components := make([][]Node, 0)
	finishOrder := g.TopologicalSort()
	for i := range finishOrder {
		finishOrder[i].node.state = unseen
	}
	for _, sink := range finishOrder {
		if g.nodes[sink.node.index].state == unseen {
			component := make([]Node, 0)
			g.dfsReversedEdges(g.nodes[sink.node.index], &component)
			components = append(components, component)
		}
	}
	return components
}

// RandMinimumCut runs Kargers algorithm to find a random minimum cut
// on the graph. If iterations is < 1, this will return an empty slice.
// Otherwise, it returns a slice of the edges crossing the best minimum
// cut found in any iteration. Call rand.Seed() before using this function.
//
// This function takes a number of iterations to start concurrently. If
// concurrent is <= 1, it will run one iteration at a time.
//
// If the graph is Directed, this will return a cut of edges in both directions.
// If the graph is Undirected, this will return a proper min cut.
func (g *Graph) RandMinimumCut(iterations, concurrent int) []Edge {
	var mutex sync.Mutex
	doneChan := make(chan struct{}, iterations)
	if concurrent < 1 {
		concurrent = 1
	}
	sem := make(chan struct{}, concurrent)
	for i := 0; i < concurrent; i++ {
		sem <- struct{}{}
	}
	// make a lite slice of the edges
	var baseAllEdges []lite.Edge
	for n := range g.nodes {
		for _, edge := range g.nodes[n].edges {
			if g.Kind == Undirected && n < edge.end.index {
				continue
			}
			baseAllEdges = append(baseAllEdges, lite.Edge{Start: n,
				End: edge.end.index, S: g.nodes[n], E: edge})
		}
	}

	minCutLite := make([]lite.Edge, len(baseAllEdges))

	for iter := 0; iter < iterations; iter++ {
		<-sem
		go func() {
			nodecount := len(g.nodes)
			allEdges := make([]lite.Edge, len(baseAllEdges))
			copy(allEdges, baseAllEdges)
			// shuffle for random edge removal order
			shuffle(allEdges)
			for nodecount > 2 {
				// remove first edge, keep the start node, collapse the end node
				// anything that points to the collapsing node now points to the keep node
				// anything that starts at the collapsing node now starts at the keep node
				keep := allEdges[len(allEdges)-1].Start
				remove := allEdges[len(allEdges)-1].End
				allEdges = allEdges[:len(allEdges)-1]
				for e := 0; e < len(allEdges); e++ {
					if allEdges[e].Start == remove {
						allEdges[e].Start = keep
					}
					if allEdges[e].End == remove {
						allEdges[e].End = keep
					}
					// remove the edge if it self looped
					if allEdges[e].Start == allEdges[e].End {
						allEdges[e] = allEdges[len(allEdges)-1]
						allEdges = allEdges[:len(allEdges)-1]
						e--
					}
				}
				// every edge removed removes a node
				nodecount--
			}

			mutex.Lock()
			if iter == 0 || len(allEdges) < len(minCutLite) {
				minCutLite = make([]lite.Edge, len(allEdges))
				copy(minCutLite, allEdges)
			}
			mutex.Unlock()

			doneChan <- struct{}{}
			sem <- struct{}{}
		}()
	}
	for iter := 0; iter < iterations; iter++ {
		<-doneChan
	}

	minCut := make([]Edge, len(minCutLite))
	for i := range minCutLite {
		start := minCutLite[i].S.(*node)
		edge := minCutLite[i].E.(edge)
		minCut[i] = Edge{Weight: edge.weight, Start: start.container,
			End: edge.end.container}
	}
	return minCut
}

// Fischer-Yates shuffle
func shuffle(edges []lite.Edge) {
	for i := len(edges) - 1; i > 0; i-- {
		j := rand.Intn(i + 1)
		edges[j], edges[i] = edges[i], edges[j]
	}
}

// MinimumSpanningTree will return the edges corresponding to the
// minimum spanning tree in the graph based off of edge weight values.
// This will return nil for a directed graph.
func (g *Graph) MinimumSpanningTree() []Edge {
	if g.Kind == Directed {
		return nil
	}
	// create priority queue for vertices
	// node.data is the index into the heap
	// node.state is the weight of an edge to its parent
	nodesBase := nodeSlice(make([]*node, len(g.nodes)))
	copy(nodesBase, g.nodes)
	for i := range nodesBase {
		nodesBase[i].state = 1<<31 - 1
		nodesBase[i].data = i
	}
	nodesBase[0].state = 0
	nodesBase[0].parent = nil
	nodes := &nodesBase
	nodes.heapInit()

	for len(*nodes) > 0 {
		min := nodes.pop()
		for _, edge := range min.edges {
			v := edge.end // get the other side of the edge
			if nodes.heapContains(v) && edge.weight < v.state {
				v.parent = min
				nodes.update(v.data, edge.weight)
			}
		}
	}

	mst := make([]Edge, 0, len(g.nodes) - 1)
	for i := range g.nodes {
		if g.nodes[i].parent != nil {
			mst = append(mst, Edge{Weight: g.nodes[i].state,
				Start: g.nodes[i].container, End: g.nodes[i].parent.container})
		}
	}

	return mst
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"github.com/GeertJohan/go.linenoise"
	//"github.com/davecgh/go-spew/spew"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/elazarl/hadoophelpers/go/lib/hadoopconf"
	"github.com/foize/go.sgr"
	"github.com/jessevdk/go-flags"
	"github.com/elazarl/hadoophelpers/go/lib/table"
	//"github.com/wsxiaoys/terminal"
)

type getOpts struct {}

type setOpts struct {}

type envAddOpts struct {}

type envSetOpts struct {}

type envOpts struct {}

func (o getOpts) Execute(args []string) error {
	opt.executed = true
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	t := table.New(4)
	c := opt.getConf()
	keys := []string{}
	for _, key := range c.Keys() {
		for _, arg := range args {
			if ok, _ := filepath.Match(arg, key); ok {
				keys = append(keys, key)
				break
			}
		}
	}
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[1].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[2].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[3].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[3].PadRight = []byte(sgr.Reset)
	}
	for _, arg := range keys {
		v, src := c.SourceGet(arg)
		if v == "" && src == hadoopconf.NoSource {
			t.Add("", arg, "", "no property")
		} else {
			t.Add(filepath.Base(src.Source), arg, "=", v)
		}
	}
	fmt.Print(t.String())
	return nil
}

func (o setOpts) Execute(args []string) error {
	opt.executed = true
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	for _, arg := range args {
		parts := strings.SplitN(arg, "=", 2)
		if len(parts) != 2 {
			return errors.New("set accepts arguments of the form x=y, no '=' in " + arg)
		}
		opt.getConf().SetIfExist(parts[0], parts[1])
	}
	opt.getConf().Save()
	return nil
}

func (o envSetOpts) Execute(args []string) error {
	opt.executed = true
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	v := opt.getEnv().Get(args[0])
	if v == nil {
		fmt.Println("No such variable", v)
	}
	t := table.New(3)
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[1].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[2].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[2].PadRight = []byte(sgr.Reset)
	}
	t.Add(v.Name, "was", v.Val)
	fmt.Println(v.Name, "was", v.Val)
	v.Val = strings.Join(args[1:], " ")
	t.Add("", "now", v.Val)
	if err := opt.getEnv().Save(); err != nil {
		return err
	}
	fmt.Print(t.String())
	return nil
}

func (o envAddOpts) Execute(args []string) error {
	opt.executed = true
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	v := opt.getEnv().Get(args[0])
	if v == nil {
		fmt.Println("No such variable", v)
	}
	t := table.New(3)
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[1].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[2].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[2].PadRight = []byte(sgr.Reset)
	}
	t.Add(v.Name, "was", v.Val)
	fmt.Println(v.Name, "was", v.Val)
	v.Append(strings.Join(args[1:], " "))
	t.Add("", "now", v.Val)
	if err := opt.getEnv().Save(); err != nil {
		return err
	}
	fmt.Print(t.String())
	return nil
}

func (o envOpts) Execute(args []string) error {
	opt.executed = true
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	t := table.New(3)
	c := opt.getEnv()
	keys := []string{}
	for _, key := range c.Keys() {
		for _, arg := range args {
			if ok, _ := filepath.Match(arg, key); ok {
				keys = append(keys, key)
				break
			}
		}
	}
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[1].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[2].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[2].PadRight = []byte(sgr.Reset)
	}
	for _, arg := range keys {
		v := c.Get(arg)
		if v == nil {
			t.Add(arg, "", "no property")
		} else {
			t.Add(arg, "=", v.Val)
		}
	}
	fmt.Print(t.String())
	return nil
}

func (o *gOpts) UseColors() bool {
	if o.Color == "auto" {
		return IsTerminal(os.Stdout.Fd())
	}
	return o.Color == "true" || o.Color == "t" || o.Color == "1"
}

type gOpts struct {
	Get getOpts `command:"get"`
	Set setOpts `command:"set"`
	SetEnv envSetOpts `command:"envset"`
	AddEnv envAddOpts `command:"envadd"`
	Env envOpts `command:"env"`
	Verbose bool `short:"v" long:"verbose" default:"true" description:"Show verbose debug information"`
	Color string `long:"color" description:"use colors on output" default:"auto"`
	ConfPath string `short:"c" long:"conf" description:"Set hadoop configuration dir"`
	conf *hadoopconf.HadoopConf
	env hadoopconf.Envs
	executed bool
}

func (opt *gOpts) setConfPath() {
	var p = "."
	if opt.ConfPath != "" {
		p = opt.ConfPath
	} else if os.Getenv("HADOOP_CONF") != "" {
		p = os.Getenv("HADOOP_CONF")
	}
	opt.ConfPath = p
}

func (opt *gOpts) getEnv() hadoopconf.Envs {
	if opt.env != nil {
		return opt.env
	}
	var err error
	opt.setConfPath()
	opt.env, err = hadoopconf.NewEnv(opt.ConfPath)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	return opt.env
}

func (opt *gOpts) getConf() *hadoopconf.HadoopConf {
	if opt.conf != nil {
		return opt.conf
	}
	var err error
	var p = "."
	if opt.ConfPath != "" {
		p = opt.ConfPath
	} else if os.Getenv("HADOOP_CONF") != "" {
		p = os.Getenv("HADOOP_CONF")
	}
	opt.conf, err = hadoopconf.New(p)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	return opt.conf
}

var opt gOpts
var conf *hadoopconf.HadoopConf

func main() {
	parser := flags.NewParser(&opt, flags.HelpFlag + flags.PassDoubleDash)
	if _, err := parser.ParseArgs(os.Args[1:]); err != nil && opt.executed {
		fmt.Println("dead:", err)
		os.Exit(1)
	}
	opt.getConf() // make sure we have correct conf
	if !opt.executed {
		if !IsTerminal(os.Stdout.Fd()) {
			fmt.Println("terminal not recognized or not supported (windows)")
			return
		}
		for {
			str, err := linenoise.Line("hadoopconf> ")
			linenoise.AddHistory(str)
			if err != nil {
				if err != linenoise.KillSignalError {
					fmt.Println("Unexpected error: %s", err)
				}
				break
			}
			if args, err := parser.ParseArgs(strings.Fields(str)); err != nil {
				fmt.Println("error:", err)
			} else if len(args) > 0 {
				fmt.Println("excessive arguments:", args)
			}
		}
	}
}
<commit_msg>no raw printf, just table<commit_after>package main

import (
	"errors"
	"github.com/GeertJohan/go.linenoise"
	//"github.com/davecgh/go-spew/spew"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/elazarl/hadoophelpers/go/lib/hadoopconf"
	"github.com/foize/go.sgr"
	"github.com/jessevdk/go-flags"
	"github.com/elazarl/hadoophelpers/go/lib/table"
	//"github.com/wsxiaoys/terminal"
)

type getOpts struct {}

type setOpts struct {}

type envAddOpts struct {}

type envSetOpts struct {}

type envOpts struct {}

func (o getOpts) Execute(args []string) error {
	opt.executed = true
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	t := table.New(4)
	c := opt.getConf()
	keys := []string{}
	for _, key := range c.Keys() {
		for _, arg := range args {
			if ok, _ := filepath.Match(arg, key); ok {
				keys = append(keys, key)
				break
			}
		}
	}
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[1].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[2].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[3].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[3].PadRight = []byte(sgr.Reset)
	}
	for _, arg := range keys {
		v, src := c.SourceGet(arg)
		if v == "" && src == hadoopconf.NoSource {
			t.Add("", arg, "", "no property")
		} else {
			t.Add(filepath.Base(src.Source), arg, "=", v)
		}
	}
	fmt.Print(t.String())
	return nil
}

func (o setOpts) Execute(args []string) error {
	opt.executed = true
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	for _, arg := range args {
		parts := strings.SplitN(arg, "=", 2)
		if len(parts) != 2 {
			return errors.New("set accepts arguments of the form x=y, no '=' in " + arg)
		}
		opt.getConf().SetIfExist(parts[0], parts[1])
	}
	opt.getConf().Save()
	return nil
}

func (o envSetOpts) Execute(args []string) error {
	opt.executed = true
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	v := opt.getEnv().Get(args[0])
	if v == nil {
		fmt.Println("No such variable", v)
	}
	t := table.New(3)
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[1].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[2].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[2].PadRight = []byte(sgr.Reset)
	}
	t.Add(v.Name, "was", v.Val)
	v.Val = strings.Join(args[1:], " ")
	t.Add("", "now", v.Val)
	if err := opt.getEnv().Save(); err != nil {
		return err
	}
	fmt.Print(t.String())
	return nil
}

func (o envAddOpts) Execute(args []string) error {
	opt.executed = true
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	v := opt.getEnv().Get(args[0])
	if v == nil {
		fmt.Println("No such variable", v)
	}
	t := table.New(3)
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[1].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[2].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[2].PadRight = []byte(sgr.Reset)
	}
	t.Add(v.Name, "was", v.Val)
	v.Append(strings.Join(args[1:], " "))
	t.Add("", "now", v.Val)
	if err := opt.getEnv().Save(); err != nil {
		return err
	}
	fmt.Print(t.String())
	return nil
}

func (o envOpts) Execute(args []string) error {
	opt.executed = true
	if len(args) == 0 {
		return errors.New("get must have nonzero number arguments")
	}
	t := table.New(3)
	c := opt.getEnv()
	keys := []string{}
	for _, key := range c.Keys() {
		for _, arg := range args {
			if ok, _ := filepath.Match(arg, key); ok {
				keys = append(keys, key)
				break
			}
		}
	}
	if opt.UseColors() {
		t.CellConf[0].PadLeft = []byte(sgr.FgCyan)
		t.CellConf[1].PadLeft = []byte(sgr.FgGrey)
		t.CellConf[2].PadLeft = []byte(sgr.ResetForegroundColor + sgr.Bold)
		t.CellConf[2].PadRight = []byte(sgr.Reset)
	}
	for _, arg := range keys {
		v := c.Get(arg)
		if v == nil {
			t.Add(arg, "", "no property")
		} else {
			t.Add(arg, "=", v.Val)
		}
	}
	fmt.Print(t.String())
	return nil
}

func (o *gOpts) UseColors() bool {
	if o.Color == "auto" {
		return IsTerminal(os.Stdout.Fd())
	}
	return o.Color == "true" || o.Color == "t" || o.Color == "1"
}

type gOpts struct {
	Get getOpts `command:"get"`
	Set setOpts `command:"set"`
	SetEnv envSetOpts `command:"envset"`
	AddEnv envAddOpts `command:"envadd"`
	Env envOpts `command:"env"`
	Verbose bool `short:"v" long:"verbose" default:"true" description:"Show verbose debug information"`
	Color string `long:"color" description:"use colors on output" default:"auto"`
	ConfPath string `short:"c" long:"conf" description:"Set hadoop configuration dir"`
	conf *hadoopconf.HadoopConf
	env hadoopconf.Envs
	executed bool
}

func (opt *gOpts) setConfPath() {
	var p = "."
	if opt.ConfPath != "" {
		p = opt.ConfPath
	} else if os.Getenv("HADOOP_CONF") != "" {
		p = os.Getenv("HADOOP_CONF")
	}
	opt.ConfPath = p
}

func (opt *gOpts) getEnv() hadoopconf.Envs {
	if opt.env != nil {
		return opt.env
	}
	var err error
	opt.setConfPath()
	opt.env, err = hadoopconf.NewEnv(opt.ConfPath)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	return opt.env
}

func (opt *gOpts) getConf() *hadoopconf.HadoopConf {
	if opt.conf != nil {
		return opt.conf
	}
	var err error
	var p = "."
	if opt.ConfPath != "" {
		p = opt.ConfPath
	} else if os.Getenv("HADOOP_CONF") != "" {
		p = os.Getenv("HADOOP_CONF")
	}
	opt.conf, err = hadoopconf.New(p)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	return opt.conf
}

var opt gOpts
var conf *hadoopconf.HadoopConf

func main() {
	parser := flags.NewParser(&opt, flags.HelpFlag + flags.PassDoubleDash)
	if _, err := parser.ParseArgs(os.Args[1:]); err != nil && opt.executed {
		fmt.Println("dead:", err)
		os.Exit(1)
	}
	opt.getConf() // make sure we have correct conf
	if !opt.executed {
		if !IsTerminal(os.Stdout.Fd()) {
			fmt.Println("terminal not recognized or not supported (windows)")
			return
		}
		for {
			str, err := linenoise.Line("hadoopconf> ")
			linenoise.AddHistory(str)
			if err != nil {
				if err != linenoise.KillSignalError {
					fmt.Println("Unexpected error: %s", err)
				}
				break
			}
			if args, err := parser.ParseArgs(strings.Fields(str)); err != nil {
				fmt.Println("error:", err)
			} else if len(args) > 0 {
				fmt.Println("excessive arguments:", args)
			}
		}
	}
}
<|endoftext|>