<commit_before>package container

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"

	"errors"
	"strings"

	"github.com/cloudfoundry-incubator/garden"
	"github.com/pivotal-cf-experimental/garden-windows/process"

	"code.google.com/p/go.net/websocket"
	"github.com/pivotal-golang/lager"
)

type container struct {
	containerizerURL url.URL
	handle           string
	logger           lager.Logger
}

type netInResponse struct {
	HostPort    uint32 `json:"hostPort"`
	ErrorString string `json:"error"`
}

type ProcessStreamEvent struct {
	MessageType    string             `json:"type"`
	ApiProcessSpec garden.ProcessSpec `json:"pspec"`
	Data           string             `json:"data"`
}

func NewContainer(containerizerURL url.URL, handle string, logger lager.Logger) *container {
	return &container{
		containerizerURL: containerizerURL,
		handle:           handle,
		logger:           logger,
	}
}

var ErrReadFromPath = errors.New("could not read tar path")

func (container *container) Handle() string {
	return container.handle
}

func (container *container) Stop(kill bool) error {
	requestUrl := container.containerizerURL
	requestUrl.Path += "/api/containers/" + container.Handle() + "/stop"
	response, err := http.Post(requestUrl.String(), "text/plain", nil)
	if err != nil {
		return err
	}
	if response.StatusCode != http.StatusOK {
		return errors.New(response.Status)
	}
	return nil
}

func (container *container) Info() (garden.ContainerInfo, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/properties"
	container.logger.Info("GETTING INFO", lager.Data{
		"url": url,
	})
	response, err := http.Get(url)
	if err != nil {
		container.logger.Info("ERROR GETTING PROPERTIES", lager.Data{
			"error": err,
		})
		return garden.ContainerInfo{}, err
	}
	defer response.Body.Close()
	rawJSON, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return garden.ContainerInfo{}, err
	}
	properties := garden.Properties{}
	err = json.Unmarshal(rawJSON, &properties)
	if err != nil {
		container.logger.Info("ERROR UNMARSHALING PROPERTIES", lager.Data{
			"error": err,
		})
		return garden.ContainerInfo{}, nil
	}
	mappedPorts := garden.PortMapping{
		HostPort:      8080,
		ContainerPort: 8080,
	}
	containerInfo := garden.ContainerInfo{
		Properties:  properties,
		HostIP:      "HOST_IP",
		ContainerIP: "CONTAINER_IP",
		ExternalIP:  container.containerizerHost(),
		MappedPorts: []garden.PortMapping{mappedPorts},
	}
	container.logger.Info("CONTAINER INFO", lager.Data{
		"containerInfo": containerInfo,
	})

	return containerInfo, nil
}

func (container *container) StreamIn(dstPath string, tarStream io.Reader) error {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/files?destination=" + dstPath

	req, err := http.NewRequest("PUT", url, tarStream)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	resp.Body.Close()
	return nil
}

func (container *container) StreamOut(srcPath string) (io.ReadCloser, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/files?source=" + srcPath
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		return nil, ErrReadFromPath
	}
	return resp.Body, nil
}

func (container *container) LimitBandwidth(limits garden.BandwidthLimits) error {
	return nil
}
func (container *container) CurrentBandwidthLimits() (garden.BandwidthLimits, error) {
	return garden.BandwidthLimits{}, nil
}

func (container *container) LimitCPU(limits garden.CPULimits) error {
	return nil
}

func (container *container) CurrentCPULimits() (garden.CPULimits, error) {
	return garden.CPULimits{}, nil
}

func (container *container) LimitDisk(limits garden.DiskLimits) error {
	return nil
}
func (container *container) CurrentDiskLimits() (garden.DiskLimits, error) {
	return garden.DiskLimits{}, nil
}

func (container *container) LimitMemory(limits garden.MemoryLimits) error {
	return nil
}
func (container *container) CurrentMemoryLimits() (garden.MemoryLimits, error) {
	return garden.MemoryLimits{}, nil
}

func (container *container) NetIn(hostPort, containerPort uint32) (uint32, uint32, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/net/in"
	response, err := http.Post(url, "application/json", strings.NewReader(fmt.Sprintf(`{"hostPort": %v}`, hostPort)))
	if err != nil {
		return 0, 0, err
	}
	responseBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return 0, 0, err
	}

	var responseJSON netInResponse
	err = json.Unmarshal(responseBody, &responseJSON)
	if err != nil {
		return 0, 0, err
	}

	if responseJSON.ErrorString != "" {
		return 0, 0, errors.New(responseJSON.ErrorString)
	}

	return responseJSON.HostPort, containerPort, err
}

func (container *container) NetOut(network string, port uint32, portRange string, protocol garden.Protocol, icmpType int32, icmpCode int32, log bool) error {
	return nil
}

func (container *container) containerizerWS() string {
	u2 := container.containerizerURL
	u2.Scheme = "ws"
	return u2.String()
}

func (container *container) Run(processSpec garden.ProcessSpec, processIO garden.ProcessIO) (garden.Process, error) {
	origin := "http://localhost/"
	wsUri := container.containerizerWS() + "/api/containers/" + container.handle + "/run"
	ws, err := websocket.Dial(wsUri, "", origin)
	if err != nil {
		return nil, err
	}
	websocket.JSON.Send(ws, ProcessStreamEvent{
		MessageType:    "run",
		ApiProcessSpec: processSpec,
	})

	proc := process.NewDotNetProcess()

	streamWebsocketIOToContainerizer(ws, processIO)
	go func() {
		err := streamWebsocketIOFromContainerizer(ws, processIO)
		if err != nil {
			proc.StreamOpen <- err.Error()
		}
		close(proc.StreamOpen)
	}()

	return proc, nil
}

func (container *container) Attach(uint32, garden.ProcessIO) (garden.Process, error) {
	return process.NewDotNetProcess(), nil
}

func (container *container) GetProperty(name string) (string, error) {
	requestUrl := container.containerizerURL
	requestUrl.Path += "/api/containers/" + container.Handle() + "/properties/" + name
	container.logger.Info("GET PROPERTY", lager.Data{
		"property": name,
		"url":      requestUrl,
	})
	response, err := http.Get(requestUrl.String())
	defer response.Body.Close()

	if err != nil {
		container.logger.Info("NO GET FOR YOU!", lager.Data{
			"error": err,
		})
		return "", err
	}
	property, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", err
	}
	return string(property), nil
}

func (container *container) SetProperty(name string, value string) error {
	requestUrl := container.containerizerURL
	requestUrl.Path += "/api/containers/" + container.Handle() + "/properties/" + name
	container.logger.Info("SET PROPERTY", lager.Data{
		"property": name,
		"url":      requestUrl,
	})
	request, err := http.NewRequest("PUT", requestUrl.String(), strings.NewReader(value))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return err
	}

	response.Body.Close()
	return nil
}

func (container *container) RemoveProperty(name string) error {
	requestUrl := container.containerizerURL
	requestUrl.Path += "/api/containers/" + container.Handle() + "/properties/" + name
	container.logger.Info("REMOVING PROPERTY", lager.Data{
		"property": name,
		"url":      requestUrl,
	})
	request, err := http.NewRequest("DELETE", requestUrl.String(), strings.NewReader(""))
	if err != nil {
		return err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return err
	}

	response.Body.Close()
	return nil
}

func (container *container) containerizerHost() string {
	return strings.Split(container.containerizerURL.Host, ":")[0]
}

func streamWebsocketIOToContainerizer(ws *websocket.Conn, processIO garden.ProcessIO) {
	if processIO.Stdin != nil {
		fiw := faninWriter{
			hasSink: make(chan struct{}),
		}
		fiw.AddSink(ws)
		fiw.AddSource(processIO.Stdin)
	}
}

func streamWebsocketIOFromContainerizer(ws *websocket.Conn, processIO garden.ProcessIO) error {
	receiveStream := ProcessStreamEvent{}
	for {
		err := websocket.JSON.Receive(ws, &receiveStream)
		if err != nil {
			return err
		}

		if receiveStream.MessageType == "stdout" && processIO.Stdout != nil {
			io.WriteString(processIO.Stdout, receiveStream.Data)
		}
		if receiveStream.MessageType == "stderr" && processIO.Stderr != nil {
			io.WriteString(processIO.Stderr, receiveStream.Data)
		}

		if receiveStream.MessageType == "error" {
			return errors.New(receiveStream.Data)
		}
		if receiveStream.MessageType == "close" {
			return nil
		}
	}
}
<commit_msg>Revert "Add "log bool" to Container#NetOut"<commit_after>package container

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"

	"errors"
	"strings"

	"github.com/cloudfoundry-incubator/garden"
	"github.com/pivotal-cf-experimental/garden-windows/process"

	"code.google.com/p/go.net/websocket"
	"github.com/pivotal-golang/lager"
)

type container struct {
	containerizerURL url.URL
	handle           string
	logger           lager.Logger
}

type netInResponse struct {
	HostPort    uint32 `json:"hostPort"`
	ErrorString string `json:"error"`
}

type ProcessStreamEvent struct {
	MessageType    string             `json:"type"`
	ApiProcessSpec garden.ProcessSpec `json:"pspec"`
	Data           string             `json:"data"`
}

func NewContainer(containerizerURL url.URL, handle string, logger lager.Logger) *container {
	return &container{
		containerizerURL: containerizerURL,
		handle:           handle,
		logger:           logger,
	}
}

var ErrReadFromPath = errors.New("could not read tar path")

func (container *container) Handle() string {
	return container.handle
}

func (container *container) Stop(kill bool) error {
	requestUrl := container.containerizerURL
	requestUrl.Path += "/api/containers/" + container.Handle() + "/stop"
	response, err := http.Post(requestUrl.String(), "text/plain", nil)
	if err != nil {
		return err
	}
	if response.StatusCode != http.StatusOK {
		return errors.New(response.Status)
	}
	return nil
}

func (container *container) Info() (garden.ContainerInfo, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/properties"
	container.logger.Info("GETTING INFO", lager.Data{
		"url": url,
	})
	response, err := http.Get(url)
	if err != nil {
		container.logger.Info("ERROR GETTING PROPERTIES", lager.Data{
			"error": err,
		})
		return garden.ContainerInfo{}, err
	}
	defer response.Body.Close()
	rawJSON, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return garden.ContainerInfo{}, err
	}
	properties := garden.Properties{}
	err = json.Unmarshal(rawJSON, &properties)
	if err != nil {
		container.logger.Info("ERROR UNMARSHALING PROPERTIES", lager.Data{
			"error": err,
		})
		return garden.ContainerInfo{}, nil
	}
	mappedPorts := garden.PortMapping{
		HostPort:      8080,
		ContainerPort: 8080,
	}
	containerInfo := garden.ContainerInfo{
		Properties:  properties,
		HostIP:      "HOST_IP",
		ContainerIP: "CONTAINER_IP",
		ExternalIP:  container.containerizerHost(),
		MappedPorts: []garden.PortMapping{mappedPorts},
	}
	container.logger.Info("CONTAINER INFO", lager.Data{
		"containerInfo": containerInfo,
	})

	return containerInfo, nil
}

func (container *container) StreamIn(dstPath string, tarStream io.Reader) error {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/files?destination=" + dstPath

	req, err := http.NewRequest("PUT", url, tarStream)
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}

	resp.Body.Close()
	return nil
}

func (container *container) StreamOut(srcPath string) (io.ReadCloser, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/files?source=" + srcPath
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		return nil, ErrReadFromPath
	}
	return resp.Body, nil
}

func (container *container) LimitBandwidth(limits garden.BandwidthLimits) error {
	return nil
}
func (container *container) CurrentBandwidthLimits() (garden.BandwidthLimits, error) {
	return garden.BandwidthLimits{}, nil
}

func (container *container) LimitCPU(limits garden.CPULimits) error {
	return nil
}

func (container *container) CurrentCPULimits() (garden.CPULimits, error) {
	return garden.CPULimits{}, nil
}

func (container *container) LimitDisk(limits garden.DiskLimits) error {
	return nil
}
func (container *container) CurrentDiskLimits() (garden.DiskLimits, error) {
	return garden.DiskLimits{}, nil
}

func (container *container) LimitMemory(limits garden.MemoryLimits) error {
	return nil
}
func (container *container) CurrentMemoryLimits() (garden.MemoryLimits, error) {
	return garden.MemoryLimits{}, nil
}

func (container *container) NetIn(hostPort, containerPort uint32) (uint32, uint32, error) {
	url := container.containerizerURL.String() + "/api/containers/" + container.Handle() + "/net/in"
	response, err := http.Post(url, "application/json", strings.NewReader(fmt.Sprintf(`{"hostPort": %v}`, hostPort)))
	if err != nil {
		return 0, 0, err
	}
	responseBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return 0, 0, err
	}

	var responseJSON netInResponse
	err = json.Unmarshal(responseBody, &responseJSON)
	if err != nil {
		return 0, 0, err
	}

	if responseJSON.ErrorString != "" {
		return 0, 0, errors.New(responseJSON.ErrorString)
	}

	return responseJSON.HostPort, containerPort, err
}

func (container *container) NetOut(network string, port uint32, portRange string, protocol garden.Protocol, icmpType int32, icmpCode int32) error {
	return nil
}

func (container *container) containerizerWS() string {
	u2 := container.containerizerURL
	u2.Scheme = "ws"
	return u2.String()
}

func (container *container) Run(processSpec garden.ProcessSpec, processIO garden.ProcessIO) (garden.Process, error) {
	origin := "http://localhost/"
	wsUri := container.containerizerWS() + "/api/containers/" + container.handle + "/run"
	ws, err := websocket.Dial(wsUri, "", origin)
	if err != nil {
		return nil, err
	}
	websocket.JSON.Send(ws, ProcessStreamEvent{
		MessageType:    "run",
		ApiProcessSpec: processSpec,
	})

	proc := process.NewDotNetProcess()

	streamWebsocketIOToContainerizer(ws, processIO)
	go func() {
		err := streamWebsocketIOFromContainerizer(ws, processIO)
		if err != nil {
			proc.StreamOpen <- err.Error()
		}
		close(proc.StreamOpen)
	}()

	return proc, nil
}

func (container *container) Attach(uint32, garden.ProcessIO) (garden.Process, error) {
	return process.NewDotNetProcess(), nil
}

func (container *container) GetProperty(name string) (string, error) {
	requestUrl := container.containerizerURL
	requestUrl.Path += "/api/containers/" + container.Handle() + "/properties/" + name
	container.logger.Info("GET PROPERTY", lager.Data{
		"property": name,
		"url":      requestUrl,
	})
	response, err := http.Get(requestUrl.String())
	defer response.Body.Close()

	if err != nil {
		container.logger.Info("NO GET FOR YOU!", lager.Data{
			"error": err,
		})
		return "", err
	}
	property, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return "", err
	}
	return string(property), nil
}

func (container *container) SetProperty(name string, value string) error {
	requestUrl := container.containerizerURL
	requestUrl.Path += "/api/containers/" + container.Handle() + "/properties/" + name
	container.logger.Info("SET PROPERTY", lager.Data{
		"property": name,
		"url":      requestUrl,
	})
	request, err := http.NewRequest("PUT", requestUrl.String(), strings.NewReader(value))
	if err != nil {
		return err
	}
	request.Header.Set("Content-Type", "application/json")
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return err
	}

	response.Body.Close()
	return nil
}

func (container *container) RemoveProperty(name string) error {
	requestUrl := container.containerizerURL
	requestUrl.Path += "/api/containers/" + container.Handle() + "/properties/" + name
	container.logger.Info("REMOVING PROPERTY", lager.Data{
		"property": name,
		"url":      requestUrl,
	})
	request, err := http.NewRequest("DELETE", requestUrl.String(), strings.NewReader(""))
	if err != nil {
		return err
	}
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return err
	}

	response.Body.Close()
	return nil
}

func (container *container) containerizerHost() string {
	return strings.Split(container.containerizerURL.Host, ":")[0]
}

func streamWebsocketIOToContainerizer(ws *websocket.Conn, processIO garden.ProcessIO) {
	if processIO.Stdin != nil {
		fiw := faninWriter{
			hasSink: make(chan struct{}),
		}
		fiw.AddSink(ws)
		fiw.AddSource(processIO.Stdin)
	}
}

func streamWebsocketIOFromContainerizer(ws *websocket.Conn, processIO garden.ProcessIO) error {
	receiveStream := ProcessStreamEvent{}
	for {
		err := websocket.JSON.Receive(ws, &receiveStream)
		if err != nil {
			return err
		}

		if receiveStream.MessageType == "stdout" && processIO.Stdout != nil {
			io.WriteString(processIO.Stdout, receiveStream.Data)
		}
		if receiveStream.MessageType == "stderr" && processIO.Stderr != nil {
			io.WriteString(processIO.Stderr, receiveStream.Data)
		}

		if receiveStream.MessageType == "error" {
			return errors.New(receiveStream.Data)
		}
		if receiveStream.MessageType == "close" {
			return nil
		}
	}
}
<|endoftext|>


<commit_before>package vm

import "testing"

func TestNewBool(t *testing.T) {
	for _, b := range []bool{true, false} {
		if bool(NewBool(b).Eval().(Bool)) != b {
			t.Fail()
		}
	}
}

func TestAnd(t *testing.T) {
	and := func(ts ...*Thunk) bool {
		return bool(And(ts...).Eval().(Bool))
	}

	if !and(True, True) {
		t.Fail()
	}

	for _, ts := range [][]*Thunk{{False, False}, {True, False}, {False, True}} {
		if and(ts...) {
			t.Fail()
		}
	}
}

func TestOr(t *testing.T) {
	or := func(ts ...*Thunk) bool {
		return bool(Or(ts...).Eval().(Bool))
	}

	for _, ts := range [][]*Thunk{{True, True}, {True, False}, {False, True}} {
		if !or(ts...) {
			t.Fail()
		}
	}

	if or(False, False) {
		t.Fail()
	}
}

func TestNot(t *testing.T) {
	for _, b := range []bool{true, false} {
		if bool(Not(NewBool(b)).Eval().(Bool)) == b {
			t.Fail()
		}
	}
}
<commit_msg>Test Equality of Bools<commit_after>package vm

import "testing"

func TestNewBool(t *testing.T) {
	for _, b := range []bool{true, false} {
		if bool(NewBool(b).Eval().(Bool)) != b {
			t.Fail()
		}
	}
}

func TestAnd(t *testing.T) {
	and := func(ts ...*Thunk) bool {
		return bool(And(ts...).Eval().(Bool))
	}

	if !and(True, True) {
		t.Fail()
	}

	for _, ts := range [][]*Thunk{{False, False}, {True, False}, {False, True}} {
		if and(ts...) {
			t.Fail()
		}
	}
}

func TestOr(t *testing.T) {
	or := func(ts ...*Thunk) bool {
		return bool(Or(ts...).Eval().(Bool))
	}

	for _, ts := range [][]*Thunk{{True, True}, {True, False}, {False, True}} {
		if !or(ts...) {
			t.Fail()
		}
	}

	if or(False, False) {
		t.Fail()
	}
}

func TestNot(t *testing.T) {
	for _, b := range []bool{true, false} {
		if bool(Not(NewBool(b)).Eval().(Bool)) == b {
			t.Fail()
		}
	}
}

func TestBoolEqual(t *testing.T) {
	equal := func(ts ...*Thunk) bool {
		return bool(Equal(ts...).Eval().(Bool))
	}

	for _, ts := range [][]*Thunk{
		{True},
		{False},
		{True, True},
		{False, False},
		{True, True, True},
	} {
		if !equal(ts...) {
			t.Fail()
		}
	}

	for _, ts := range [][]*Thunk{
		{True, False},
		{False, True},
		{True, True, False},
	} {
		if equal(ts...) {
			t.Fail()
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 The StudyGolang Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// http://studygolang.com
// Author:polaris	polaris@studygolang.com

package logic

import (
	. "db"
	"errors"
	"fmt"
	"io/ioutil"
	"model"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
	"github.com/polaris1119/config"
	"github.com/polaris1119/logger"
	"github.com/tidwall/gjson"
	"golang.org/x/net/context"
)

var titlePattern = config.ConfigFile.MustValue("crawl", "article_title_pattern")

type AutoCrawlLogic struct{}

var DefaultAutoCrawl = AutoCrawlLogic{}

func (self AutoCrawlLogic) DoCrawl(isAll bool) error {
	autoCrawlConfList := make([]*model.AutoCrawlRule, 0)
	err := MasterDB.Where("status=?", model.AutoCrawlOn).Find(&autoCrawlConfList)
	if err != nil {
		logger.Errorln("AutoCrawlLogic FindBy Error:", err)
		return err
	}

	for _, autoCrawlConf := range autoCrawlConfList {
		self.crawlOneWebsite(autoCrawlConf, isAll)
	}

	return nil
}

// 通过网站标识抓取
func (self AutoCrawlLogic) CrawlWebsite(website string, isAll bool) error {
	autoCrawlConf := &model.AutoCrawlRule{}
	_, err := MasterDB.Where("website=?", website).Get(autoCrawlConf)
	if err != nil {
		return err
	}

	if autoCrawlConf.Id == 0 {
		return errors.New("the website is not exists in auto crawl rule.")
	}

	go self.crawlOneWebsite(autoCrawlConf, isAll)

	return nil
}

func (self AutoCrawlLogic) crawlOneWebsite(autoCrawlConf *model.AutoCrawlRule, isAll bool) {
	maxPage := 1
	crawlUrl := autoCrawlConf.IncrUrl
	if isAll {
		crawlUrl = autoCrawlConf.AllUrl
		maxPage = autoCrawlConf.MaxPage
	}

	pageField := autoCrawlConf.PageField

	// 个人博客，一般通过 tag 方式获取，这种处理方式和搜索不一样
	if autoCrawlConf.Keywords == "" {
		for p := maxPage; p >= 1; p-- {
			curUrl := ""

			if pageField == "" {
				if p > 1 {
					curUrl += crawlUrl + "page/" + strconv.Itoa(p)
				} else {
					curUrl = crawlUrl
				}
			} else {
				page := fmt.Sprintf("?%s=%d", pageField, p)
				curUrl += crawlUrl + page
			}

			// 标题不包含 go 等关键词的，也入库
			if err := self.parseArticleList(curUrl, autoCrawlConf, false); err != nil {
				logger.Errorln("parse article url", curUrl, "error:", err)
				break
			}
		}
		return
	}

	keywords := strings.Split(autoCrawlConf.Keywords, ",")
	for _, keyword := range keywords {
		for p := 1; p <= maxPage; p++ {

			curUrl := ""
			page := fmt.Sprintf("&%s=%d", pageField, p)
			if strings.Contains(crawlUrl, "%s") {
				curUrl = fmt.Sprintf(crawlUrl, keyword) + page
			} else {
				curUrl = crawlUrl + keyword + page
			}

			var err error
			if _, ok := autoCrawlConf.ExtMap["json_api"]; ok {
				err = self.fetchArticleListFromApi(curUrl, autoCrawlConf, true)
			} else {
				err = self.parseArticleList(curUrl, autoCrawlConf, true)
			}
			if err != nil {
				logger.Errorln("parse article url", curUrl, "error:", err)
			}

			time.Sleep(30 * time.Second)
		}
	}
}

func (self AutoCrawlLogic) parseArticleList(strUrl string, autoCrawlConf *model.AutoCrawlRule, isSearch bool) (err error) {

	logger.Infoln("parse url:", strUrl)

	var doc *goquery.Document

	if autoCrawlConf.ExtMap == nil {
		doc, err = goquery.NewDocument(strUrl)
	} else {
		var req *http.Request
		req, err = http.NewRequest("GET", strUrl, nil)
		if err != nil {
			return err
		}
		if referer, ok := autoCrawlConf.ExtMap["referer"]; ok {
			req.Header.Add("Referer", referer)
		}

		var resp *http.Response
		resp, err = http.DefaultClient.Do(req)
		if err != nil {
			return err
		}
		doc, err = goquery.NewDocumentFromResponse(resp)
	}

	if err != nil {
		return
	}

	listSelector := autoCrawlConf.ListSelector
	resultSelector := autoCrawlConf.ResultSelector

	u, err := url.Parse(autoCrawlConf.IncrUrl)
	if err != nil {
		logger.Errorln("parse incr_url error:", err)
		return
	}
	host := u.Scheme + "://" + u.Host

	articleSelection := doc.Find(listSelector)
	// 后面的先入库
	for i := articleSelection.Length() - 1; i >= 0; i-- {

		contentSelection := goquery.NewDocumentFromNode(articleSelection.Get(i)).Selection

		aSelection := contentSelection.Find(resultSelector)

		// 搜索时，避免搜到垃圾，对标题进一步判断
		if isSearch && titlePattern != "" {
			title := aSelection.Text()

			matched, err := regexp.MatchString(titlePattern, title)
			if err != nil {
				logger.Errorln(err)
				continue
			}

			if !matched {
				continue
			}
		}

		articleUrl, ok := aSelection.Attr("href")
		if ok {
			pos := strings.LastIndex(articleUrl, "?")
			if pos != -1 {
				articleUrl = articleUrl[:pos]
			}

			if !strings.HasPrefix(articleUrl, "http") {
				articleUrl = host + articleUrl
			}
			DefaultArticle.ParseArticle(context.Background(), articleUrl, isSearch)
		}
	}

	return
}

func (self AutoCrawlLogic) fetchArticleListFromApi(strUrl string, autoCrawlConf *model.AutoCrawlRule, isSearch bool) error {
	logger.Infoln("parse api url:", strUrl)

	req, err := http.NewRequest("GET", strUrl, nil)
	if err != nil {
		return err
	}
	req.Header.Add("accept", "application/json")
	req.Header.Add("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	u, err := url.Parse(autoCrawlConf.IncrUrl)
	if err != nil {
		logger.Errorln("parse incr_url error:", err)
		return err
	}
	host := u.Scheme + "://" + u.Host

	result := gjson.ParseBytes(body)
	result = result.Get(autoCrawlConf.ListSelector)
	result.ForEach(func(key, value gjson.Result) bool {
		articleUrl := value.Get(autoCrawlConf.ResultSelector).String()

		pos := strings.LastIndex(articleUrl, "?")
		if pos != -1 {
			articleUrl = articleUrl[:pos]
		}

		if strings.HasPrefix(articleUrl, "/") {
			articleUrl = host + articleUrl
		} else if !strings.HasPrefix(articleUrl, "http") {
			// jianshu 写死
			articleUrl = host + "/p/" + articleUrl
		}
		DefaultArticle.ParseArticle(context.Background(), articleUrl, isSearch)

		return true
	})

	return nil
}
<commit_msg>简书是 POST 方式<commit_after>// Copyright 2017 The StudyGolang Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// http://studygolang.com
// Author:polaris	polaris@studygolang.com

package logic

import (
	. "db"
	"errors"
	"fmt"
	"io/ioutil"
	"model"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
	"github.com/polaris1119/config"
	"github.com/polaris1119/logger"
	"github.com/tidwall/gjson"
	"golang.org/x/net/context"
)

var titlePattern = config.ConfigFile.MustValue("crawl", "article_title_pattern")

type AutoCrawlLogic struct{}

var DefaultAutoCrawl = AutoCrawlLogic{}

func (self AutoCrawlLogic) DoCrawl(isAll bool) error {
	autoCrawlConfList := make([]*model.AutoCrawlRule, 0)
	err := MasterDB.Where("status=?", model.AutoCrawlOn).Find(&autoCrawlConfList)
	if err != nil {
		logger.Errorln("AutoCrawlLogic FindBy Error:", err)
		return err
	}

	for _, autoCrawlConf := range autoCrawlConfList {
		self.crawlOneWebsite(autoCrawlConf, isAll)
	}

	return nil
}

// 通过网站标识抓取
func (self AutoCrawlLogic) CrawlWebsite(website string, isAll bool) error {
	autoCrawlConf := &model.AutoCrawlRule{}
	_, err := MasterDB.Where("website=?", website).Get(autoCrawlConf)
	if err != nil {
		return err
	}

	if autoCrawlConf.Id == 0 {
		return errors.New("the website is not exists in auto crawl rule.")
	}

	go self.crawlOneWebsite(autoCrawlConf, isAll)

	return nil
}

func (self AutoCrawlLogic) crawlOneWebsite(autoCrawlConf *model.AutoCrawlRule, isAll bool) {
	maxPage := 1
	crawlUrl := autoCrawlConf.IncrUrl
	if isAll {
		crawlUrl = autoCrawlConf.AllUrl
		maxPage = autoCrawlConf.MaxPage
	}

	pageField := autoCrawlConf.PageField

	// 个人博客，一般通过 tag 方式获取，这种处理方式和搜索不一样
	if autoCrawlConf.Keywords == "" {
		for p := maxPage; p >= 1; p-- {
			curUrl := ""

			if pageField == "" {
				if p > 1 {
					curUrl += crawlUrl + "page/" + strconv.Itoa(p)
				} else {
					curUrl = crawlUrl
				}
			} else {
				page := fmt.Sprintf("?%s=%d", pageField, p)
				curUrl += crawlUrl + page
			}

			// 标题不包含 go 等关键词的，也入库
			if err := self.parseArticleList(curUrl, autoCrawlConf, false); err != nil {
				logger.Errorln("parse article url", curUrl, "error:", err)
				break
			}
		}
		return
	}

	keywords := strings.Split(autoCrawlConf.Keywords, ",")
	for _, keyword := range keywords {
		for p := 1; p <= maxPage; p++ {

			curUrl := ""
			page := fmt.Sprintf("&%s=%d", pageField, p)
			if strings.Contains(crawlUrl, "%s") {
				curUrl = fmt.Sprintf(crawlUrl, keyword) + page
			} else {
				curUrl = crawlUrl + keyword + page
			}

			var err error
			if _, ok := autoCrawlConf.ExtMap["json_api"]; ok {
				err = self.fetchArticleListFromApi(curUrl, autoCrawlConf, true)
			} else {
				err = self.parseArticleList(curUrl, autoCrawlConf, true)
			}
			if err != nil {
				logger.Errorln("parse article url", curUrl, "error:", err)
			}

			time.Sleep(30 * time.Second)
		}
	}
}

func (self AutoCrawlLogic) parseArticleList(strUrl string, autoCrawlConf *model.AutoCrawlRule, isSearch bool) (err error) {

	logger.Infoln("parse url:", strUrl)

	var doc *goquery.Document

	if autoCrawlConf.ExtMap == nil {
		doc, err = goquery.NewDocument(strUrl)
	} else {
		var req *http.Request
		req, err = http.NewRequest("GET", strUrl, nil)
		if err != nil {
			return err
		}
		if referer, ok := autoCrawlConf.ExtMap["referer"]; ok {
			req.Header.Add("Referer", referer)
		}

		var resp *http.Response
		resp, err = http.DefaultClient.Do(req)
		if err != nil {
			return err
		}
		doc, err = goquery.NewDocumentFromResponse(resp)
	}

	if err != nil {
		return
	}

	listSelector := autoCrawlConf.ListSelector
	resultSelector := autoCrawlConf.ResultSelector

	u, err := url.Parse(autoCrawlConf.IncrUrl)
	if err != nil {
		logger.Errorln("parse incr_url error:", err)
		return
	}
	host := u.Scheme + "://" + u.Host

	articleSelection := doc.Find(listSelector)
	// 后面的先入库
	for i := articleSelection.Length() - 1; i >= 0; i-- {

		contentSelection := goquery.NewDocumentFromNode(articleSelection.Get(i)).Selection

		aSelection := contentSelection.Find(resultSelector)

		// 搜索时，避免搜到垃圾，对标题进一步判断
		if isSearch && titlePattern != "" {
			title := aSelection.Text()

			matched, err := regexp.MatchString(titlePattern, title)
			if err != nil {
				logger.Errorln(err)
				continue
			}

			if !matched {
				continue
			}
		}

		articleUrl, ok := aSelection.Attr("href")
		if ok {
			pos := strings.LastIndex(articleUrl, "?")
			if pos != -1 {
				articleUrl = articleUrl[:pos]
			}

			if !strings.HasPrefix(articleUrl, "http") {
				articleUrl = host + articleUrl
			}
			DefaultArticle.ParseArticle(context.Background(), articleUrl, isSearch)
		}
	}

	return
}

func (self AutoCrawlLogic) fetchArticleListFromApi(strUrl string, autoCrawlConf *model.AutoCrawlRule, isSearch bool) error {
	logger.Infoln("parse api url:", strUrl)

	// jianshu must be POST
	req, err := http.NewRequest("POST", strUrl, nil)
	if err != nil {
		return err
	}
	req.Header.Add("accept", "application/json")
	req.Header.Add("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/57.0.2987.133 Safari/537.36")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}

	u, err := url.Parse(autoCrawlConf.IncrUrl)
	if err != nil {
		logger.Errorln("parse incr_url error:", err)
		return err
	}
	host := u.Scheme + "://" + u.Host

	result := gjson.ParseBytes(body)
	result = result.Get(autoCrawlConf.ListSelector)
	result.ForEach(func(key, value gjson.Result) bool {
		articleUrl := value.Get(autoCrawlConf.ResultSelector).String()

		pos := strings.LastIndex(articleUrl, "?")
		if pos != -1 {
			articleUrl = articleUrl[:pos]
		}

		if strings.HasPrefix(articleUrl, "/") {
			articleUrl = host + articleUrl
		} else if !strings.HasPrefix(articleUrl, "http") {
			// jianshu 写死
			articleUrl = host + "/p/" + articleUrl
		}
		DefaultArticle.ParseArticle(context.Background(), articleUrl, isSearch)

		return true
	})

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 The Serviced 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.

// +build unit

package facade_test

import (
	"time"

	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/domain/host"
	"github.com/control-center/serviced/domain/pool"
	"github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/utils"
	"github.com/stretchr/testify/mock"
	. "gopkg.in/check.v1"
)

func (ft *FacadeUnitTest) Test_GetResourcePool(c *C) {
	ft.setupMockDFSLocking()
	poolID := "somePoolID"
	expectedPool := pool.ResourcePool{ID: poolID}
	key := pool.Key(poolID)
	ft.hostStore.On("FindHostsWithPoolID", ft.ctx, mock.AnythingOfType("string")).
		Return([]host.Host{
			{ID: "host1", Cores: 2, Memory: 3},
			{ID: "host2", Cores: 3, Memory: 4},
		}, nil)
	ft.poolStore.On("Get", ft.ctx, key, mock.AnythingOfType("*pool.ResourcePool")).
		Return(nil).
		Run(func(args mock.Arguments) {
			pool := args.Get(2).(*pool.ResourcePool)
			*pool = expectedPool
		})

	result, err := ft.Facade.GetResourcePool(ft.ctx, poolID)

	c.Assert(err, IsNil)
	c.Assert(result.ID, Equals, poolID)
	c.Assert(result.CoreCapacity, Equals, 5)
	c.Assert(result.MemoryCapacity, Equals, uint64(7))
}

func (ft *FacadeUnitTest) Test_GetResourcePoolWithNoHosts(c *C) {
	ft.setupMockDFSLocking()
	poolID := "somePoolID"
	expectedPool := pool.ResourcePool{ID: poolID}
	key := pool.Key(poolID)
	ft.hostStore.On("FindHostsWithPoolID", ft.ctx, mock.AnythingOfType("string")).
		Return([]host.Host{}, nil)
	ft.poolStore.On("Get", ft.ctx, key, mock.AnythingOfType("*pool.ResourcePool")).
		Return(nil).
		Run(func(args mock.Arguments) {
			pool := args.Get(2).(*pool.ResourcePool)
			*pool = expectedPool
		})

	result, err := ft.Facade.GetResourcePool(ft.ctx, poolID)

	c.Assert(err, IsNil)
	c.Assert(result, Not(IsNil))
	c.Assert(result.ID, Equals, poolID)
	c.Assert(result.CoreCapacity, Equals, 0)
	c.Assert(result.MemoryCapacity, Equals, uint64(0))
}

func (ft *FacadeUnitTest) Test_GetResourcePoolFailsForNoSuchEntity(c *C) {
	ft.setupMockDFSLocking()
	poolID := "somePoolID"
	key := pool.Key(poolID)

	ft.poolStore.On("Get", ft.ctx, key, mock.AnythingOfType("*pool.ResourcePool")).Return(datastore.ErrNoSuchEntity{})

	result, err := ft.Facade.GetResourcePool(ft.ctx, poolID)

	c.Assert(result, IsNil)
	c.Assert(err, IsNil)
}

func (ft *FacadeUnitTest) Test_GetResourcePoolFailsForOtherDBError(c *C) {
	ft.setupMockDFSLocking()
	poolID := "somePoolID"
	key := pool.Key(poolID)
	expectedError := datastore.ErrEmptyKind

	ft.poolStore.On("Get", ft.ctx, key, mock.AnythingOfType("*pool.ResourcePool")).Return(expectedError)

	result, err := ft.Facade.GetResourcePool(ft.ctx, poolID)

	c.Assert(result, IsNil)
	c.Assert(err, Equals, expectedError)
}

func (ft *FacadeUnitTest) Test_GetReadPoolsShouldReturnCorrectValues(c *C) {
	ft.setupMockDFSLocking()

	resourcePool := pool.ResourcePool{
		ID:                "firstPool",
		Description:       "The first pool",
		MemoryCapacity:    0,
		MemoryCommitment:  0,
		CoreCapacity:      0,
		ConnectionTimeout: 10,
		CreatedAt:         time.Now(),
		UpdatedAt:         time.Now(),
		Permissions:       pool.DFSAccess,
	}

	firstHost := host.Host{
		ID:     "firstHost",
		Cores:  6,
		Memory: 12000,
	}

	secondHost := host.Host{
		ID:     "secondHost",
		Cores:  8,
		Memory: 10000,
	}

	firstService := service.Service{
		ID: "firstService",
		RAMCommitment: utils.EngNotation{
			Value: uint64(1000),
		},
	}

	secondService := service.Service{
		ID: "secondService",
		RAMCommitment: utils.EngNotation{
			Value: uint64(2000),
		},
	}

	ft.hostStore.On("FindHostsWithPoolID", ft.ctx, resourcePool.ID).
		Return([]host.Host{firstHost, secondHost}, nil)

	ft.poolStore.On("GetResourcePools", ft.ctx).
		Return([]pool.ResourcePool{resourcePool}, nil)

	ft.serviceStore.On("GetServicesByPool", ft.ctx, resourcePool.ID).
		Return([]service.Service{firstService, secondService}, nil)

	pools, err := ft.Facade.GetReadPools(ft.ctx)
	c.Assert(err, IsNil)
	c.Assert(pools, Not(IsNil))
	c.Assert(len(pools), Equals, 1)

	p := pools[0]
	c.Assert(p.ID, Equals, resourcePool.ID)
	c.Assert(p.CoreCapacity, Equals, 14)
	c.Assert(p.MemoryCapacity, Equals, uint64(22000))
	c.Assert(p.MemoryCommitment, Equals, uint64(3000))
	c.Assert(p.ConnectionTimeout, Equals, 10)
	c.Assert(p.CreatedAt, TimeEqual, resourcePool.CreatedAt)
	c.Assert(p.UpdatedAt, TimeEqual, resourcePool.UpdatedAt)
	c.Assert(p.DFSAccess, Equals, resourcePool.Permissions&pool.DFSAccess != 0)
	c.Assert(p.AdminAccess, Equals, resourcePool.Permissions&pool.AdminAccess != 0)
}
<commit_msg>update pool unit test<commit_after>// Copyright 2016 The Serviced 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.

// +build unit

package facade_test

import (
	"time"

	"github.com/control-center/serviced/datastore"
	"github.com/control-center/serviced/domain/host"
	"github.com/control-center/serviced/domain/pool"
	"github.com/control-center/serviced/domain/service"
	"github.com/control-center/serviced/utils"
	"github.com/stretchr/testify/mock"
	. "gopkg.in/check.v1"
)

func (ft *FacadeUnitTest) Test_GetResourcePool(c *C) {
	ft.setupMockDFSLocking()
	poolID := "somePoolID"
	expectedPool := pool.ResourcePool{ID: poolID}
	key := pool.Key(poolID)
	ft.hostStore.On("FindHostsWithPoolID", ft.ctx, mock.AnythingOfType("string")).
		Return([]host.Host{
			{ID: "host1", Cores: 2, Memory: 3},
			{ID: "host2", Cores: 3, Memory: 4},
		}, nil)
	ft.poolStore.On("Get", ft.ctx, key, mock.AnythingOfType("*pool.ResourcePool")).
		Return(nil).
		Run(func(args mock.Arguments) {
			pool := args.Get(2).(*pool.ResourcePool)
			*pool = expectedPool
		})

	result, err := ft.Facade.GetResourcePool(ft.ctx, poolID)

	c.Assert(err, IsNil)
	c.Assert(result.ID, Equals, poolID)
	c.Assert(result.CoreCapacity, Equals, 5)
	c.Assert(result.MemoryCapacity, Equals, uint64(7))
}

func (ft *FacadeUnitTest) Test_GetResourcePoolWithNoHosts(c *C) {
	ft.setupMockDFSLocking()
	poolID := "somePoolID"
	expectedPool := pool.ResourcePool{ID: poolID}
	key := pool.Key(poolID)
	ft.hostStore.On("FindHostsWithPoolID", ft.ctx, mock.AnythingOfType("string")).
		Return([]host.Host{}, nil)
	ft.poolStore.On("Get", ft.ctx, key, mock.AnythingOfType("*pool.ResourcePool")).
		Return(nil).
		Run(func(args mock.Arguments) {
			pool := args.Get(2).(*pool.ResourcePool)
			*pool = expectedPool
		})

	result, err := ft.Facade.GetResourcePool(ft.ctx, poolID)

	c.Assert(err, IsNil)
	c.Assert(result, Not(IsNil))
	c.Assert(result.ID, Equals, poolID)
	c.Assert(result.CoreCapacity, Equals, 0)
	c.Assert(result.MemoryCapacity, Equals, uint64(0))
}

func (ft *FacadeUnitTest) Test_GetResourcePoolFailsForNoSuchEntity(c *C) {
	ft.setupMockDFSLocking()
	poolID := "somePoolID"
	key := pool.Key(poolID)

	ft.poolStore.On("Get", ft.ctx, key, mock.AnythingOfType("*pool.ResourcePool")).Return(datastore.ErrNoSuchEntity{})

	result, err := ft.Facade.GetResourcePool(ft.ctx, poolID)

	c.Assert(result, IsNil)
	c.Assert(err, IsNil)
}

func (ft *FacadeUnitTest) Test_GetResourcePoolFailsForOtherDBError(c *C) {
	ft.setupMockDFSLocking()
	poolID := "somePoolID"
	key := pool.Key(poolID)
	expectedError := datastore.ErrEmptyKind

	ft.poolStore.On("Get", ft.ctx, key, mock.AnythingOfType("*pool.ResourcePool")).Return(expectedError)

	result, err := ft.Facade.GetResourcePool(ft.ctx, poolID)

	c.Assert(result, IsNil)
	c.Assert(err, Equals, expectedError)
}

func (ft *FacadeUnitTest) Test_GetReadPoolsShouldReturnCorrectValues(c *C) {
	ft.setupMockDFSLocking()

	resourcePool := pool.ResourcePool{
		ID:                "firstPool",
		Description:       "The first pool",
		MemoryCapacity:    0,
		MemoryCommitment:  0,
		CoreCapacity:      0,
		ConnectionTimeout: 10,
		CreatedAt:         time.Now(),
		UpdatedAt:         time.Now(),
		Permissions:       pool.DFSAccess,
	}

	firstHost := host.Host{
		ID:     "firstHost",
		Cores:  6,
		Memory: 12000,
	}

	secondHost := host.Host{
		ID:     "secondHost",
		Cores:  8,
		Memory: 10000,
	}

	firstService := service.Service{
		ID: "firstService",
		RAMCommitment: utils.EngNotation{
			Value: uint64(1000),
		},
	}

	secondService := service.Service{
		ID: "secondService",
		RAMCommitment: utils.EngNotation{
			Value: uint64(2000),
		},
	}

	ft.hostStore.On("FindHostsWithPoolID", ft.ctx, resourcePool.ID).
		Return([]host.Host{firstHost, secondHost}, nil)

	ft.poolStore.On("GetResourcePools", ft.ctx).
		Return([]pool.ResourcePool{resourcePool}, nil)

	ft.serviceStore.On("GetServicesByPool", ft.ctx, resourcePool.ID).
		Return([]service.Service{firstService, secondService}, nil)

	pools, err := ft.Facade.GetReadPools(ft.ctx)
	c.Assert(err, IsNil)
	c.Assert(pools, Not(IsNil))
	c.Assert(len(pools), Equals, 1)

	p := pools[0]
	c.Assert(p.ID, Equals, resourcePool.ID)
	c.Assert(p.CoreCapacity, Equals, 14)
	c.Assert(p.MemoryCapacity, Equals, uint64(22000))
	c.Assert(p.MemoryCommitment, Equals, uint64(3000))
	c.Assert(p.ConnectionTimeout, Equals, 10)
	c.Assert(p.CreatedAt, TimeEqual, resourcePool.CreatedAt)
	c.Assert(p.UpdatedAt, TimeEqual, resourcePool.UpdatedAt)
	c.Assert(p.Permissions, Equals, resourcePool.Permissions)
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Rook Authors. All rights reserved.

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

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

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

package crash

import (
	"context"
	"time"

	"github.com/pkg/errors"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mgr"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mon"
	"github.com/rook/rook/pkg/operator/ceph/cluster/rbd"
	appsv1 "k8s.io/api/apps/v1"
	kerrors "k8s.io/apimachinery/pkg/api/errors"

	"github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/ceph/file/mds"
	"github.com/rook/rook/pkg/operator/ceph/object"
	"github.com/rook/rook/pkg/operator/ceph/version"

	"github.com/coreos/pkg/capnslog"

	"github.com/rook/rook/pkg/operator/ceph/cluster/osd"
	"github.com/rook/rook/pkg/operator/k8sutil"

	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"

	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	"github.com/rook/rook/pkg/operator/ceph/disruption/controllerconfig"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	getVersionRetryInterval = 5
	getVersionMaxRetries    = 60
)

var (
	logger = capnslog.NewPackageLogger("github.com/rook/rook", controllerName)
	// Implement reconcile.Reconciler so the controller can reconcile objects
	_ reconcile.Reconciler = &ReconcileNode{}
)

// ReconcileNode reconciles ReplicaSets
type ReconcileNode struct {
	// client can be used to retrieve objects from the APIServer.
	scheme *runtime.Scheme
	client client.Client
}

// Reconcile reconciles a node and ensures that it has a crashcollector deployment
// attached to it.
// The Controller will requeue the Request to be processed again if an error is non-nil or
// Result.Requeue is true, otherwise upon completion it will remove the work from the queue.
func (r *ReconcileNode) Reconcile(request reconcile.Request) (reconcile.Result, error) {
	// workaround because the rook logging mechanism is not compatible with the controller-runtime logging interface
	result, err := r.reconcile(request)
	if err != nil {
		logger.Error(err)
	}
	return result, err
}

func (r *ReconcileNode) reconcile(request reconcile.Request) (reconcile.Result, error) {

	logger.Debugf("reconciling node: %q", request.Name)

	// get the node object
	node := &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: request.Name}}
	err := r.client.Get(context.TODO(), request.NamespacedName, node)
	if err != nil {
		return reconcile.Result{}, errors.Errorf("could not get node %q", request.NamespacedName)
	}

	// Get the list of all the Ceph pods
	cephPods, err := r.cephPodList()
	if err != nil {
		if len(cephPods) == 0 {
			return reconcile.Result{}, nil
		}
		return reconcile.Result{}, errors.Wrapf(err, "failed to list all ceph pods")
	}

	namespaceToPodList := make(map[string][]corev1.Pod)
	for _, cephPod := range cephPods {
		podNamespace := cephPod.GetNamespace()
		podList, ok := namespaceToPodList[podNamespace]
		if !ok {
			// initialize list
			namespaceToPodList[podNamespace] = []corev1.Pod{cephPod}
		} else {
			// append cephPod to namespace's pod list
			namespaceToPodList[podNamespace] = append(podList, cephPod)
		}
	}

	for namespace, cephPods := range namespaceToPodList {
		// get dataDirHostPath from the CephCluster
		cephClusters := &cephv1.CephClusterList{}
		err := r.client.List(context.TODO(), cephClusters, client.InNamespace(namespace))
		if err != nil {
			return reconcile.Result{}, errors.Wrapf(err, "could not get cephcluster in namespaces %q", namespace)
		}
		if len(cephClusters.Items) < 1 {
			logger.Debugf("no CephCluster found in the namespace %q", namespace)
			return reconcile.Result{}, nil
		}

		cephCluster := cephClusters.Items[0]
		if len(cephClusters.Items) > 1 {
			logger.Errorf("more than one CephCluster found in the namespace %q, choosing the first one %q", namespace, cephCluster.GetName())
		}

		// If the crash controller is disabled in the spec let's do a noop
		if cephCluster.Spec.CrashCollector.Disable {
			deploymentList := &appsv1.DeploymentList{}
			namespaceListOpts := client.InNamespace(request.Namespace)

			// Try to fetch the list of existing deployment and remove them
			err := r.client.List(context.TODO(), deploymentList, client.MatchingLabels{k8sutil.AppAttr: AppName}, namespaceListOpts)
			if err != nil {
				logger.Errorf("failed to list crash collector deployments, delete it/them manually. %v", err)
				return reconcile.Result{}, nil
			}

			//  Try to delete all the crash deployments
			for _, d := range deploymentList.Items {
				err := r.deleteCrashCollector(d)
				if err != nil {
					logger.Errorf("failed to delete crash collector deployment %q, delete it manually. %v", d.Name, err)
					continue
				}
				logger.Infof("crash collector deployment %q successfully removed", d.Name)
			}

			return reconcile.Result{}, nil
		}

		clusterImage := cephCluster.Spec.CephVersion.Image
		cephVersion, err := getImageVersion(cephCluster)
		if err != nil {
			logger.Errorf("ceph version not found for image %q used by cluster %q. %v", clusterImage, cephCluster.Name, err)
			return reconcile.Result{}, nil
		}

		uniqueTolerations := controllerconfig.TolerationSet{}
		hasCephPods := false
		for _, cephPod := range cephPods {
			if cephPod.Spec.NodeName == request.Name {
				logger.Debugf("cephPod.Spec.NodeName is %q and request.Name is %q", cephPod.Spec.NodeName, request.Name)
				hasCephPods = true
				for _, podToleration := range cephPod.Spec.Tolerations {
					// Add toleration to the map
					uniqueTolerations.Add(podToleration)
				}
			}
		}

		if hasCephPods {
			tolerations := uniqueTolerations.ToList()
			op, err := r.createOrUpdateCephCrash(*node, tolerations, cephCluster, cephVersion)
			if err != nil {
				return reconcile.Result{}, errors.Wrapf(err, "node reconcile failed on op %q", op)
			}
			logger.Debugf("deployment successfully reconciled for node %q. operation: %q", request.Name, op)
		}
	}

	return reconcile.Result{}, nil
}

func (r *ReconcileNode) cephPodList() ([]corev1.Pod, error) {
	cephPods := make([]corev1.Pod, 0)
	cephAppNames := []string{mon.AppName, mgr.AppName, osd.AppName, object.AppName, mds.AppName, rbd.AppName}

	for _, app := range cephAppNames {
		podList := &corev1.PodList{}
		err := r.client.List(context.TODO(), podList, client.MatchingLabels{k8sutil.AppAttr: app})
		if err != nil {
			return cephPods, errors.Wrapf(err, "could not list the %q pods", app)
		}

		cephPods = append(cephPods, podList.Items...)
	}

	return cephPods, nil
}

// getImageVersion returns the CephVersion registered for a specified image (if any) and whether any image was found.
func getImageVersion(cephCluster cephv1.CephCluster) (*version.CephVersion, error) {
	for i := 0; i < getVersionMaxRetries; i++ {
		// If the Ceph cluster has not yet recorded the image and version for the current image in its spec, then the Crash
		// controller should wait for the version to be detected.
		if cephCluster.Status.CephVersion != nil && cephCluster.Spec.CephVersion.Image == cephCluster.Status.CephVersion.Image {
			logger.Debugf("ceph version found %q", cephCluster.Status.CephVersion.Version)
			return controller.ExtractCephVersionFromLabel(cephCluster.Status.CephVersion.Version)
		}
		<-time.After(time.Second * getVersionRetryInterval)
	}
	return nil, errors.New("attempt to determine ceph version for the current cluster image timed out")
}

func (r *ReconcileNode) deleteCrashCollector(deployment appsv1.Deployment) error {
	deploymentName := deployment.ObjectMeta.Name
	namespace := deployment.ObjectMeta.Namespace
	dep := &appsv1.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      deploymentName,
			Namespace: namespace,
		},
	}

	err := r.client.Delete(context.TODO(), dep)
	if err != nil && !kerrors.IsNotFound(err) {
		return errors.Wrapf(err, "could not delete crash collector deployment %q", deploymentName)
	}

	return nil
}
<commit_msg>Ceph: Remove crashcollector pod automatically when the node is deleted.<commit_after>/*
Copyright 2019 The Rook Authors. All rights reserved.

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

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

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

package crash

import (
	"context"
	"time"

	"github.com/pkg/errors"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mgr"
	"github.com/rook/rook/pkg/operator/ceph/cluster/mon"
	"github.com/rook/rook/pkg/operator/ceph/cluster/rbd"
	appsv1 "k8s.io/api/apps/v1"
	kerrors "k8s.io/apimachinery/pkg/api/errors"

	"github.com/rook/rook/pkg/operator/ceph/controller"
	"github.com/rook/rook/pkg/operator/ceph/file/mds"
	"github.com/rook/rook/pkg/operator/ceph/object"
	"github.com/rook/rook/pkg/operator/ceph/version"

	"github.com/coreos/pkg/capnslog"

	"github.com/rook/rook/pkg/operator/ceph/cluster/osd"
	"github.com/rook/rook/pkg/operator/k8sutil"

	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/reconcile"

	cephv1 "github.com/rook/rook/pkg/apis/ceph.rook.io/v1"
	"github.com/rook/rook/pkg/operator/ceph/disruption/controllerconfig"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

const (
	getVersionRetryInterval = 5
	getVersionMaxRetries    = 60
)

var (
	logger = capnslog.NewPackageLogger("github.com/rook/rook", controllerName)
	// Implement reconcile.Reconciler so the controller can reconcile objects
	_ reconcile.Reconciler = &ReconcileNode{}
)

// ReconcileNode reconciles ReplicaSets
type ReconcileNode struct {
	// client can be used to retrieve objects from the APIServer.
	scheme *runtime.Scheme
	client client.Client
}

// Reconcile reconciles a node and ensures that it has a crashcollector deployment
// attached to it.
// The Controller will requeue the Request to be processed again if an error is non-nil or
// Result.Requeue is true, otherwise upon completion it will remove the work from the queue.
func (r *ReconcileNode) Reconcile(request reconcile.Request) (reconcile.Result, error) {
	// workaround because the rook logging mechanism is not compatible with the controller-runtime logging interface
	result, err := r.reconcile(request)
	if err != nil {
		logger.Error(err)
	}
	return result, err
}

func (r *ReconcileNode) reconcile(request reconcile.Request) (reconcile.Result, error) {

	logger.Debugf("reconciling node: %q", request.Name)

	// get the node object
	node := &corev1.Node{ObjectMeta: metav1.ObjectMeta{Name: request.Name}}
	err := r.client.Get(context.TODO(), request.NamespacedName, node)
	if err != nil {
		//if a node is not present, check if there are any crashcollector deployment for that node and delete it.
		deploymentList := &appsv1.DeploymentList{}
		namespaceListOpts := client.InNamespace(request.Namespace)
		err := r.client.List(context.TODO(), deploymentList, client.MatchingLabels{k8sutil.AppAttr: AppName, NodeNameLabel: request.Name}, namespaceListOpts)
		if err != nil {
			logger.Errorf("failed to list crash collector deployments, delete it/them manually. %v", err)
		}
		for _, d := range deploymentList.Items {
			logger.Infof("deleting deployment %q for deleted node %q", d.ObjectMeta.Name, request.Name)
			err := r.deleteCrashCollector(d)
			if err != nil {
				logger.Errorf("failed to delete crash collector deployment %q, delete it manually. %v", d.Name, err)
				continue
			}
			logger.Infof("crash collector deployment %q successfully removed", d.Name)
		}
		return reconcile.Result{}, errors.Errorf("could not get node %q", request.NamespacedName)
	}

	// Get the list of all the Ceph pods
	cephPods, err := r.cephPodList()
	if err != nil {
		if len(cephPods) == 0 {
			return reconcile.Result{}, nil
		}
		return reconcile.Result{}, errors.Wrapf(err, "failed to list all ceph pods")
	}

	namespaceToPodList := make(map[string][]corev1.Pod)
	for _, cephPod := range cephPods {
		podNamespace := cephPod.GetNamespace()
		podList, ok := namespaceToPodList[podNamespace]
		if !ok {
			// initialize list
			namespaceToPodList[podNamespace] = []corev1.Pod{cephPod}
		} else {
			// append cephPod to namespace's pod list
			namespaceToPodList[podNamespace] = append(podList, cephPod)
		}
	}

	for namespace, cephPods := range namespaceToPodList {
		// get dataDirHostPath from the CephCluster
		cephClusters := &cephv1.CephClusterList{}
		err := r.client.List(context.TODO(), cephClusters, client.InNamespace(namespace))
		if err != nil {
			return reconcile.Result{}, errors.Wrapf(err, "could not get cephcluster in namespaces %q", namespace)
		}
		if len(cephClusters.Items) < 1 {
			logger.Debugf("no CephCluster found in the namespace %q", namespace)
			return reconcile.Result{}, nil
		}

		cephCluster := cephClusters.Items[0]
		if len(cephClusters.Items) > 1 {
			logger.Errorf("more than one CephCluster found in the namespace %q, choosing the first one %q", namespace, cephCluster.GetName())
		}

		// If the crash controller is disabled in the spec let's do a noop
		if cephCluster.Spec.CrashCollector.Disable {
			deploymentList := &appsv1.DeploymentList{}
			namespaceListOpts := client.InNamespace(request.Namespace)

			// Try to fetch the list of existing deployment and remove them
			err := r.client.List(context.TODO(), deploymentList, client.MatchingLabels{k8sutil.AppAttr: AppName}, namespaceListOpts)
			if err != nil {
				logger.Errorf("failed to list crash collector deployments, delete it/them manually. %v", err)
				return reconcile.Result{}, nil
			}

			//  Try to delete all the crash deployments
			for _, d := range deploymentList.Items {
				err := r.deleteCrashCollector(d)
				if err != nil {
					logger.Errorf("failed to delete crash collector deployment %q, delete it manually. %v", d.Name, err)
					continue
				}
				logger.Infof("crash collector deployment %q successfully removed", d.Name)
			}

			return reconcile.Result{}, nil
		}

		clusterImage := cephCluster.Spec.CephVersion.Image
		cephVersion, err := getImageVersion(cephCluster)
		if err != nil {
			logger.Errorf("ceph version not found for image %q used by cluster %q. %v", clusterImage, cephCluster.Name, err)
			return reconcile.Result{}, nil
		}

		uniqueTolerations := controllerconfig.TolerationSet{}
		hasCephPods := false
		for _, cephPod := range cephPods {
			if cephPod.Spec.NodeName == request.Name {
				logger.Debugf("cephPod.Spec.NodeName is %q and request.Name is %q", cephPod.Spec.NodeName, request.Name)
				hasCephPods = true
				for _, podToleration := range cephPod.Spec.Tolerations {
					// Add toleration to the map
					uniqueTolerations.Add(podToleration)
				}
			}
		}

		if hasCephPods {
			tolerations := uniqueTolerations.ToList()
			op, err := r.createOrUpdateCephCrash(*node, tolerations, cephCluster, cephVersion)
			if err != nil {
				return reconcile.Result{}, errors.Wrapf(err, "node reconcile failed on op %q", op)
			}
			logger.Debugf("deployment successfully reconciled for node %q. operation: %q", request.Name, op)
		}
	}

	return reconcile.Result{}, nil
}

func (r *ReconcileNode) cephPodList() ([]corev1.Pod, error) {
	cephPods := make([]corev1.Pod, 0)
	cephAppNames := []string{mon.AppName, mgr.AppName, osd.AppName, object.AppName, mds.AppName, rbd.AppName}

	for _, app := range cephAppNames {
		podList := &corev1.PodList{}
		err := r.client.List(context.TODO(), podList, client.MatchingLabels{k8sutil.AppAttr: app})
		if err != nil {
			return cephPods, errors.Wrapf(err, "could not list the %q pods", app)
		}

		cephPods = append(cephPods, podList.Items...)
	}

	return cephPods, nil
}

// getImageVersion returns the CephVersion registered for a specified image (if any) and whether any image was found.
func getImageVersion(cephCluster cephv1.CephCluster) (*version.CephVersion, error) {
	for i := 0; i < getVersionMaxRetries; i++ {
		// If the Ceph cluster has not yet recorded the image and version for the current image in its spec, then the Crash
		// controller should wait for the version to be detected.
		if cephCluster.Status.CephVersion != nil && cephCluster.Spec.CephVersion.Image == cephCluster.Status.CephVersion.Image {
			logger.Debugf("ceph version found %q", cephCluster.Status.CephVersion.Version)
			return controller.ExtractCephVersionFromLabel(cephCluster.Status.CephVersion.Version)
		}
		<-time.After(time.Second * getVersionRetryInterval)
	}
	return nil, errors.New("attempt to determine ceph version for the current cluster image timed out")
}

func (r *ReconcileNode) deleteCrashCollector(deployment appsv1.Deployment) error {
	deploymentName := deployment.ObjectMeta.Name
	namespace := deployment.ObjectMeta.Namespace
	dep := &appsv1.Deployment{
		ObjectMeta: metav1.ObjectMeta{
			Name:      deploymentName,
			Namespace: namespace,
		},
	}

	err := r.client.Delete(context.TODO(), dep)
	if err != nil && !kerrors.IsNotFound(err) {
		return errors.Wrapf(err, "could not delete crash collector deployment %q", deploymentName)
	}

	return nil
}
<|endoftext|>


<commit_before>package jsoniter

import (
	"fmt"
	"io"
	"reflect"
	"unsafe"
)

func decoderOfArray(cfg *frozenConfig, prefix string, typ reflect.Type) ValDecoder {
	decoder := decoderOfType(cfg, prefix+"[array]->", typ.Elem())
	return &arrayDecoder{typ, typ.Elem(), decoder}
}

func encoderOfArray(cfg *frozenConfig, prefix string, typ reflect.Type) ValEncoder {
	encoder := encoderOfType(cfg, prefix+"[array]->", typ.Elem())
	if typ.Elem().Kind() == reflect.Map {
		encoder = &OptionalEncoder{encoder}
	}
	return &arrayEncoder{typ, typ.Elem(), encoder}
}

type arrayEncoder struct {
	arrayType   reflect.Type
	elemType    reflect.Type
	elemEncoder ValEncoder
}

func (encoder *arrayEncoder) Encode(ptr unsafe.Pointer, stream *Stream) {
	if encoder.arrayType.Len() == 0 {
		stream.WriteEmptyArray()
		return
	}
	stream.WriteArrayStart()
	elemPtr := unsafe.Pointer(ptr)
	encoder.elemEncoder.Encode(elemPtr, stream)
	for i := 1; i < encoder.arrayType.Len(); i++ {
		stream.WriteMore()
		elemPtr = unsafe.Pointer(uintptr(elemPtr) + encoder.elemType.Size())
		encoder.elemEncoder.Encode(unsafe.Pointer(elemPtr), stream)
	}
	stream.WriteArrayEnd()
	if stream.Error != nil && stream.Error != io.EOF {
		stream.Error = fmt.Errorf("%v: %s", encoder.arrayType, stream.Error.Error())
	}
}

func (encoder *arrayEncoder) EncodeInterface(val interface{}, stream *Stream) {
	// special optimization for interface{}
	e := (*emptyInterface)(unsafe.Pointer(&val))
	if e.word == nil {
		stream.WriteArrayStart()
		stream.WriteNil()
		stream.WriteArrayEnd()
		return
	}
	elemType := encoder.arrayType.Elem()
	if encoder.arrayType.Len() == 1 && (elemType.Kind() == reflect.Ptr || elemType.Kind() == reflect.Map) {
		ptr := uintptr(e.word)
		e.word = unsafe.Pointer(&ptr)
	}
	if reflect.TypeOf(val).Kind() == reflect.Ptr {
		encoder.Encode(unsafe.Pointer(&e.word), stream)
	} else {
		encoder.Encode(e.word, stream)
	}
}

func (encoder *arrayEncoder) IsEmpty(ptr unsafe.Pointer) bool {
	return false
}

type arrayDecoder struct {
	arrayType   reflect.Type
	elemType    reflect.Type
	elemDecoder ValDecoder
}

func (decoder *arrayDecoder) Decode(ptr unsafe.Pointer, iter *Iterator) {
	decoder.doDecode(ptr, iter)
	if iter.Error != nil && iter.Error != io.EOF {
		iter.Error = fmt.Errorf("%v: %s", decoder.arrayType, iter.Error.Error())
	}
}

func (decoder *arrayDecoder) doDecode(ptr unsafe.Pointer, iter *Iterator) {
	offset := uintptr(0)
	iter.ReadArrayCB(func(iter *Iterator) bool {
		if offset < decoder.arrayType.Size() {
			decoder.elemDecoder.Decode(unsafe.Pointer(uintptr(ptr)+offset), iter)
			offset += decoder.elemType.Size()
		} else {
			iter.Skip()
		}
		return true
	})
}
<commit_msg>Optimize 0-length array case<commit_after>package jsoniter

import (
	"fmt"
	"io"
	"reflect"
	"unsafe"
)

func decoderOfArray(cfg *frozenConfig, prefix string, typ reflect.Type) ValDecoder {
	decoder := decoderOfType(cfg, prefix+"[array]->", typ.Elem())
	return &arrayDecoder{typ, typ.Elem(), decoder}
}

func encoderOfArray(cfg *frozenConfig, prefix string, typ reflect.Type) ValEncoder {
	if typ.Len() == 0 {
		return emptyArrayEncoder{}
	}
	encoder := encoderOfType(cfg, prefix+"[array]->", typ.Elem())
	if typ.Elem().Kind() == reflect.Map {
		encoder = &OptionalEncoder{encoder}
	}
	return &arrayEncoder{typ, typ.Elem(), encoder}
}

type emptyArrayEncoder struct{}

func (encoder emptyArrayEncoder) Encode(ptr unsafe.Pointer, stream *Stream) {
	stream.WriteEmptyArray()
}

func (encoder emptyArrayEncoder) EncodeInterface(val interface{}, stream *Stream) {
	stream.WriteEmptyArray()
}

func (encoder emptyArrayEncoder) IsEmpty(ptr unsafe.Pointer) bool {
	return true
}

type arrayEncoder struct {
	arrayType   reflect.Type
	elemType    reflect.Type
	elemEncoder ValEncoder
}

func (encoder *arrayEncoder) Encode(ptr unsafe.Pointer, stream *Stream) {
	stream.WriteArrayStart()
	elemPtr := unsafe.Pointer(ptr)
	encoder.elemEncoder.Encode(elemPtr, stream)
	for i := 1; i < encoder.arrayType.Len(); i++ {
		stream.WriteMore()
		elemPtr = unsafe.Pointer(uintptr(elemPtr) + encoder.elemType.Size())
		encoder.elemEncoder.Encode(unsafe.Pointer(elemPtr), stream)
	}
	stream.WriteArrayEnd()
	if stream.Error != nil && stream.Error != io.EOF {
		stream.Error = fmt.Errorf("%v: %s", encoder.arrayType, stream.Error.Error())
	}
}

func (encoder *arrayEncoder) EncodeInterface(val interface{}, stream *Stream) {
	// special optimization for interface{}
	e := (*emptyInterface)(unsafe.Pointer(&val))
	if e.word == nil {
		stream.WriteArrayStart()
		stream.WriteNil()
		stream.WriteArrayEnd()
		return
	}
	elemType := encoder.arrayType.Elem()
	if encoder.arrayType.Len() == 1 && (elemType.Kind() == reflect.Ptr || elemType.Kind() == reflect.Map) {
		ptr := uintptr(e.word)
		e.word = unsafe.Pointer(&ptr)
	}
	if reflect.TypeOf(val).Kind() == reflect.Ptr {
		encoder.Encode(unsafe.Pointer(&e.word), stream)
	} else {
		encoder.Encode(e.word, stream)
	}
}

func (encoder *arrayEncoder) IsEmpty(ptr unsafe.Pointer) bool {
	return false
}

type arrayDecoder struct {
	arrayType   reflect.Type
	elemType    reflect.Type
	elemDecoder ValDecoder
}

func (decoder *arrayDecoder) Decode(ptr unsafe.Pointer, iter *Iterator) {
	decoder.doDecode(ptr, iter)
	if iter.Error != nil && iter.Error != io.EOF {
		iter.Error = fmt.Errorf("%v: %s", decoder.arrayType, iter.Error.Error())
	}
}

func (decoder *arrayDecoder) doDecode(ptr unsafe.Pointer, iter *Iterator) {
	offset := uintptr(0)
	iter.ReadArrayCB(func(iter *Iterator) bool {
		if offset < decoder.arrayType.Size() {
			decoder.elemDecoder.Decode(unsafe.Pointer(uintptr(ptr)+offset), iter)
			offset += decoder.elemType.Size()
		} else {
			iter.Skip()
		}
		return true
	})
}
<|endoftext|>


<commit_before>package todolist

import (
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Parser struct{}

func (p *Parser) ParseNewTodo(input string) *Todo {
	r, _ := regexp.Compile(`^(add|a)(\\ |)`)
	input = r.ReplaceAllString(input, "")
	if input == "" {
		return nil
	}

	todo := NewTodo()
	todo.Subject = p.Subject(input)
	todo.Projects = p.Projects(input)
	todo.Contexts = p.Contexts(input)
	if p.hasDue(input) {
		todo.Due = p.Due(input, time.Now())
	}
	return todo
}

func (p *Parser) Subject(input string) string {
	if strings.Contains(input, " due") {
		index := strings.LastIndex(input, " due")
		return input[0:index]
	} else {
		return input
	}
}

func (p *Parser) ExpandProject(input string) string {
	r, _ := regexp.Compile(`(ex|expand) +\d+ +\+[\p{L}\d_]+:`)
	pattern := r.FindString(input)
	if len(pattern) == 0 {
		return ""
	}

	newProject := pattern[0 : len(pattern)-1]
	project := strings.Split(newProject, " ")
	return project[len(project)-1]
}

func (p *Parser) Projects(input string) []string {
	r, _ := regexp.Compile(`\+[\p{L}\d_]+`)
	return p.matchWords(input, r)
}

func (p *Parser) Contexts(input string) []string {
	r, err := regexp.Compile(`\@[\p{L}\d_]+`)
	if err != nil {
		fmt.Println("regex error", err)
	}
	return p.matchWords(input, r)
}

func (p *Parser) hasDue(input string) bool {
	r1, _ := regexp.Compile(`due \w+$`)
	r2, _ := regexp.Compile(`due \w+ \d+$`)
	return (r1.MatchString(input) || r2.MatchString(input))
}

func (p *Parser) Due(input string, day time.Time) string {
	r, _ := regexp.Compile(`due .*$`)

	res := r.FindString(input)
	res = res[4:]
	switch res {
	case "none":
		return ""
	case "today", "tod":
		return bod(time.Now()).Format("2006-01-02")
	case "tomorrow", "tom":
		return bod(time.Now()).AddDate(0, 0, 1).Format("2006-01-02")
	case "monday", "mon":
		return p.monday(day)
	case "tuesday", "tue":
		return p.tuesday(day)
	case "wednesday", "wed":
		return p.wednesday(day)
	case "thursday", "thu":
		return p.thursday(day)
	case "friday", "fri":
		return p.friday(day)
	case "saturday", "sat":
		return p.saturday(day)
	case "sunday", "sun":
		return p.sunday(day)
	case "next week":
		n := bod(time.Now())
		return getNearestMonday(n).AddDate(0, 0, 7).Format("2006-01-02")
	}
	return p.parseArbitraryDate(res, time.Now())
}

func (p *Parser) parseArbitraryDate(_date string, pivot time.Time) string {
	d1 := p.parseArbitraryDateWithYear(_date, pivot.Year())

	var diff1 time.Duration
	if d1.After(time.Now()) {
		diff1 = d1.Sub(pivot)
	} else {
		diff1 = pivot.Sub(d1)
	}
	d2 := p.parseArbitraryDateWithYear(_date, pivot.Year()+1)
	if d2.Sub(pivot) > diff1 {
		return d1.Format("2006-01-02")
	} else {
		return d2.Format("2006-01-02")
	}
}

func (p *Parser) parseArbitraryDateWithYear(_date string, year int) time.Time {
	res := strings.Join([]string{_date, strconv.Itoa(year)}, " ")
	if date, err := time.Parse("Jan 2 2006", res); err == nil {
		return date
	}

	if date, err := time.Parse("2 Jan 2006", res); err == nil {
		return date
	}
	fmt.Printf("Could not parse the date you gave me: %s\n", _date)
	fmt.Println("I'm expecting a date like \"Dec 22\" or \"22 Dec\".")
	fmt.Println("See http://todolist.site/#adding for more info.")
	os.Exit(-1)
	return time.Now()
}

func (p *Parser) monday(day time.Time) string {
	mon := getNearestMonday(day)
	return p.thisOrNextWeek(mon, day)
}

func (p *Parser) tuesday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 1)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) wednesday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 2)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) thursday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 3)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) friday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 4)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) saturday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 5)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) sunday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 6)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) thisOrNextWeek(day time.Time, pivotDay time.Time) string {
	if day.Before(pivotDay) {
		return day.AddDate(0, 0, 7).Format("2006-01-02")
	} else {
		return day.Format("2006-01-02")
	}
}

func (p *Parser) matchWords(input string, r *regexp.Regexp) []string {
	results := r.FindAllString(input, -1)
	ret := []string{}

	for _, val := range results {
		ret = append(ret, val[1:])
	}
	return ret
}
<commit_msg>Fix spacing bug when adding todos<commit_after>package todolist

import (
	"fmt"
	"os"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type Parser struct{}

func (p *Parser) ParseNewTodo(input string) *Todo {
	r, _ := regexp.Compile(`^(add|a)(\\ |) `)
	input = r.ReplaceAllString(input, "")
	if input == "" {
		return nil
	}

	todo := NewTodo()
	todo.Subject = p.Subject(input)
	todo.Projects = p.Projects(input)
	todo.Contexts = p.Contexts(input)
	if p.hasDue(input) {
		todo.Due = p.Due(input, time.Now())
	}
	return todo
}

func (p *Parser) Subject(input string) string {
	if strings.Contains(input, " due") {
		index := strings.LastIndex(input, " due")
		return input[0:index]
	} else {
		return input
	}
}

func (p *Parser) ExpandProject(input string) string {
	r, _ := regexp.Compile(`(ex|expand) +\d+ +\+[\p{L}\d_]+:`)
	pattern := r.FindString(input)
	if len(pattern) == 0 {
		return ""
	}

	newProject := pattern[0 : len(pattern)-1]
	project := strings.Split(newProject, " ")
	return project[len(project)-1]
}

func (p *Parser) Projects(input string) []string {
	r, _ := regexp.Compile(`\+[\p{L}\d_]+`)
	return p.matchWords(input, r)
}

func (p *Parser) Contexts(input string) []string {
	r, err := regexp.Compile(`\@[\p{L}\d_]+`)
	if err != nil {
		fmt.Println("regex error", err)
	}
	return p.matchWords(input, r)
}

func (p *Parser) hasDue(input string) bool {
	r1, _ := regexp.Compile(`due \w+$`)
	r2, _ := regexp.Compile(`due \w+ \d+$`)
	return (r1.MatchString(input) || r2.MatchString(input))
}

func (p *Parser) Due(input string, day time.Time) string {
	r, _ := regexp.Compile(`due .*$`)

	res := r.FindString(input)
	res = res[4:]
	switch res {
	case "none":
		return ""
	case "today", "tod":
		return bod(time.Now()).Format("2006-01-02")
	case "tomorrow", "tom":
		return bod(time.Now()).AddDate(0, 0, 1).Format("2006-01-02")
	case "monday", "mon":
		return p.monday(day)
	case "tuesday", "tue":
		return p.tuesday(day)
	case "wednesday", "wed":
		return p.wednesday(day)
	case "thursday", "thu":
		return p.thursday(day)
	case "friday", "fri":
		return p.friday(day)
	case "saturday", "sat":
		return p.saturday(day)
	case "sunday", "sun":
		return p.sunday(day)
	case "next week":
		n := bod(time.Now())
		return getNearestMonday(n).AddDate(0, 0, 7).Format("2006-01-02")
	}
	return p.parseArbitraryDate(res, time.Now())
}

func (p *Parser) parseArbitraryDate(_date string, pivot time.Time) string {
	d1 := p.parseArbitraryDateWithYear(_date, pivot.Year())

	var diff1 time.Duration
	if d1.After(time.Now()) {
		diff1 = d1.Sub(pivot)
	} else {
		diff1 = pivot.Sub(d1)
	}
	d2 := p.parseArbitraryDateWithYear(_date, pivot.Year()+1)
	if d2.Sub(pivot) > diff1 {
		return d1.Format("2006-01-02")
	} else {
		return d2.Format("2006-01-02")
	}
}

func (p *Parser) parseArbitraryDateWithYear(_date string, year int) time.Time {
	res := strings.Join([]string{_date, strconv.Itoa(year)}, " ")
	if date, err := time.Parse("Jan 2 2006", res); err == nil {
		return date
	}

	if date, err := time.Parse("2 Jan 2006", res); err == nil {
		return date
	}
	fmt.Printf("Could not parse the date you gave me: %s\n", _date)
	fmt.Println("I'm expecting a date like \"Dec 22\" or \"22 Dec\".")
	fmt.Println("See http://todolist.site/#adding for more info.")
	os.Exit(-1)
	return time.Now()
}

func (p *Parser) monday(day time.Time) string {
	mon := getNearestMonday(day)
	return p.thisOrNextWeek(mon, day)
}

func (p *Parser) tuesday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 1)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) wednesday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 2)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) thursday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 3)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) friday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 4)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) saturday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 5)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) sunday(day time.Time) string {
	tue := getNearestMonday(day).AddDate(0, 0, 6)
	return p.thisOrNextWeek(tue, day)
}

func (p *Parser) thisOrNextWeek(day time.Time, pivotDay time.Time) string {
	if day.Before(pivotDay) {
		return day.AddDate(0, 0, 7).Format("2006-01-02")
	} else {
		return day.Format("2006-01-02")
	}
}

func (p *Parser) matchWords(input string, r *regexp.Regexp) []string {
	results := r.FindAllString(input, -1)
	ret := []string{}

	for _, val := range results {
		ret = append(ret, val[1:])
	}
	return ret
}
<|endoftext|>


<commit_before>// Package torrent provides a torrent client implementation.
package torrent

import (
	"encoding/hex"
	"errors"
	"io"
	"math/rand"
	"net"
	"net/url"
	"time"

	"github.com/cenkalti/rain/internal/clientversion"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/resume"
	"github.com/cenkalti/rain/storage"
	"github.com/cenkalti/rain/storage/filestorage"
	"github.com/cenkalti/rain/torrent/internal/acceptor"
	"github.com/cenkalti/rain/torrent/internal/addrlist"
	"github.com/cenkalti/rain/torrent/internal/allocator"
	"github.com/cenkalti/rain/torrent/internal/announcer"
	"github.com/cenkalti/rain/torrent/internal/bitfield"
	"github.com/cenkalti/rain/torrent/internal/handshaker/incominghandshaker"
	"github.com/cenkalti/rain/torrent/internal/handshaker/outgoinghandshaker"
	"github.com/cenkalti/rain/torrent/internal/infodownloader"
	"github.com/cenkalti/rain/torrent/internal/magnet"
	"github.com/cenkalti/rain/torrent/internal/metainfo"
	"github.com/cenkalti/rain/torrent/internal/mse"
	"github.com/cenkalti/rain/torrent/internal/peer"
	"github.com/cenkalti/rain/torrent/internal/piece"
	"github.com/cenkalti/rain/torrent/internal/piecedownloader"
	"github.com/cenkalti/rain/torrent/internal/torrentdata"
	"github.com/cenkalti/rain/torrent/internal/tracker"
	"github.com/cenkalti/rain/torrent/internal/tracker/httptracker"
	"github.com/cenkalti/rain/torrent/internal/tracker/udptracker"
	"github.com/cenkalti/rain/torrent/internal/verifier"
)

const (
	parallelInfoDownloads  = 4
	parallelPieceDownloads = 4
	maxPeerDial            = 40
	maxPeerAccept          = 40
)

var (
	// http://www.bittorrent.org/beps/bep_0020.html
	peerIDPrefix = []byte("-RN" + clientversion.Version + "-")
)

// Torrent connects to peers and downloads files from swarm.
type Torrent struct {
	*downloadSpec

	// Unique peer ID is generated per downloader.
	peerID [20]byte

	// Data provides IO access to pieces in torrent.
	data *torrentdata.Data

	// Contains state about the pieces in torrent.
	pieces []*piece.Piece

	// Contains pieces in sorted order for piece selection function.
	sortedPieces []*piece.Piece

	// Peers are sent to this channel when they are disconnected.
	peerDisconnectedC chan *peer.Peer

	// All messages coming from peers are sent to this channel.
	messages chan peer.Message

	// We keep connected peers in this map after they complete handshake phase.
	peers         map[*peer.Peer]struct{}
	incomingPeers map[*peer.Peer]struct{}
	outgoingPeers map[*peer.Peer]struct{}

	// Active piece downloads are kept in this map.
	pieceDownloaders map[*peer.Peer]*piecedownloader.PieceDownloader

	// Active metadata downloads are kept in this map.
	infoDownloaders map[*peer.Peer]*infodownloader.InfoDownloader

	// A peer is optimistically unchoked regardless of their download rate.
	optimisticUnchokedPeer *peer.Peer

	// This channel is closed once all pieces are downloaded and verified.
	completeC chan struct{}

	// If any unrecoverable error occurs, it will be sent to this channel and download will be stopped.
	errC chan error

	// When Stop() is called, it will close this channel to signal run() function to stop.
	closeC chan struct{}

	// This channel will be closed after run loop exists.
	doneC chan struct{}

	// These are the channels for sending a message to run() loop.
	statsCommandC       chan statsRequest       // Stats()
	startCommandC       chan struct{}           // Start()
	stopCommandC        chan struct{}           // Stop()
	notifyErrorCommandC chan notifyErrorCommand // NotifyError()

	// Trackers send announce responses to this channel.
	addrsFromTrackers chan []*net.TCPAddr

	// Keeps a list of peer addresses to connect.
	addrList *addrlist.AddrList

	// New raw connections created by OutgoingHandshaker are sent to here.
	incomingConnC chan net.Conn

	// Keep a set of peer IDs to block duplicate connections.
	peerIDs map[[20]byte]struct{}

	// Listens for incoming peer connections.
	acceptor *acceptor.Acceptor

	// Special hash of info hash for encypted connection handshake.
	sKeyHash [20]byte

	// Tracker implementations for giving to announcers.
	trackersInstances []tracker.Tracker

	// Announces the status of torrent to trackers to get peer addresses.
	announcers []*announcer.Announcer

	// List of peers in handshake state.
	incomingHandshakers map[string]*incominghandshaker.IncomingHandshaker
	outgoingHandshakers map[string]*outgoinghandshaker.OutgoingHandshaker

	// Handshake results are sent to these channels by handshakers.
	incomingHandshakerResultC chan incominghandshaker.Result
	outgoingHandshakerResultC chan outgoinghandshaker.Result

	// When metadata of the torrent downloaded completely, a message is sent to this channel.
	infoDownloaderResultC chan infodownloader.Result

	// When a piece is downloaded completely a message is sent to this channel.
	pieceDownloaderResultC chan piecedownloader.Result

	// Announcers send a request to this channel to get information about the torrent.
	announcerRequestC chan *announcer.Request

	// A timer that ticks periodically to keep a certain number of peers unchoked.
	unchokeTimer  *time.Ticker
	unchokeTimerC <-chan time.Time

	// A timer that ticks periodically to keep a random peer unchoked regardless of its upload rate.
	optimisticUnchokeTimer  *time.Ticker
	optimisticUnchokeTimerC <-chan time.Time

	allocator          *allocator.Allocator
	allocatorProgressC chan allocator.Progress
	allocatorResultC   chan allocator.Result

	verifier          *verifier.Verifier
	verifierProgressC chan verifier.Progress
	verifierResultC   chan verifier.Result

	log logger.Logger
}

// New returns a new torrent by reading a torrent metainfo file.
func New(r io.Reader, port int, sto storage.Storage) (*Torrent, error) {
	m, err := metainfo.New(r)
	if err != nil {
		return nil, err
	}
	spec := &downloadSpec{
		infoHash: m.Info.Hash,
		trackers: m.GetTrackers(),
		name:     m.Info.Name,
		port:     port,
		storage:  sto,
		info:     m.Info,
	}
	return newTorrent(spec)
}

// NewMagnet returns a new torrent by parsing a magnet link.
func NewMagnet(link string, port int, sto storage.Storage) (*Torrent, error) {
	m, err := magnet.New(link)
	if err != nil {
		return nil, err
	}
	spec := &downloadSpec{
		infoHash: m.InfoHash,
		trackers: m.Trackers,
		name:     m.Name,
		port:     port,
		storage:  sto,
	}
	return newTorrent(spec)
}

// NewResume returns a new torrent by loading all info from a resume.DB.
func NewResume(res resume.DB) (*Torrent, error) {
	spec, err := res.Read()
	if err != nil {
		return nil, err
	}
	if spec == nil {
		return nil, errors.New("no resume info")
	}
	return loadResumeSpec(spec, res)
}

// Name of the torrent.
// For magnet downloads name can change after metadata is downloaded but this method still returns the initial name.
func (t *Torrent) Name() string {
	return t.name
}

// InfoHash string encoded in hex.
// InfoHash is a unique value that identifies the files in torrent.
func (t *Torrent) InfoHash() string {
	return hex.EncodeToString(t.infoHash[:])
}

// SetResume adds resume capability to the torrent.
// It must be called before Start() is called.
func (t *Torrent) SetResume(res resume.DB) error {
	spec, err := res.Read()
	if err != nil {
		return err
	}
	if spec == nil {
		return t.writeResume(res)
	}
	t2, err := loadResumeSpec(spec, res)
	if err != nil {
		return err
	}
	if t.InfoHash() != t2.InfoHash() {
		t2.Close()
		return errors.New("invalid resume file (info hashes does not match)")
	}
	t.Close()
	*t = *t2
	return nil
}

func loadResumeSpec(spec *resume.Spec, res resume.DB) (*Torrent, error) {
	var err error
	dspec := &downloadSpec{
		port:     spec.Port,
		trackers: spec.Trackers,
		name:     spec.Name,
		resume:   res,
	}
	copy(dspec.infoHash[:], spec.InfoHash)
	if len(spec.Info) > 0 {
		dspec.info, err = metainfo.NewInfo(spec.Info)
		if err != nil {
			return nil, err
		}
		if len(spec.Bitfield) > 0 {
			dspec.bitfield = bitfield.New(dspec.info.NumPieces)
			copy(dspec.bitfield.Bytes(), spec.Bitfield)
		}
	}
	switch spec.StorageType {
	case filestorage.StorageType:
		dspec.storage = &filestorage.FileStorage{}
	default:
		return nil, errors.New("unknown storage type: " + spec.StorageType)
	}
	err = dspec.storage.Load(spec.StorageArgs)
	if err != nil {
		return nil, err
	}
	return newTorrent(dspec)
}

func (t *Torrent) writeResume(res resume.DB) error {
	rspec := &resume.Spec{
		InfoHash:    t.infoHash[:],
		Port:        t.port,
		Name:        t.name,
		Trackers:    t.trackers,
		StorageType: t.storage.Type(),
		StorageArgs: t.storage.Args(),
	}
	if t.info != nil {
		rspec.Info = t.info.Bytes
	}
	if t.bitfield != nil {
		rspec.Bitfield = t.bitfield.Bytes()
	}
	return res.Write(rspec)
}

func newTorrent(spec *downloadSpec) (*Torrent, error) {
	logName := spec.name
	if len(logName) > 8 {
		logName = logName[:8]
	}
	d := &Torrent{
		downloadSpec:              spec,
		log:                       logger.New("download " + logName),
		peerDisconnectedC:         make(chan *peer.Peer),
		messages:                  make(chan peer.Message),
		peers:                     make(map[*peer.Peer]struct{}),
		incomingPeers:             make(map[*peer.Peer]struct{}),
		outgoingPeers:             make(map[*peer.Peer]struct{}),
		pieceDownloaders:          make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		infoDownloaders:           make(map[*peer.Peer]*infodownloader.InfoDownloader),
		completeC:                 make(chan struct{}),
		closeC:                    make(chan struct{}),
		doneC:                     make(chan struct{}),
		startCommandC:             make(chan struct{}),
		stopCommandC:              make(chan struct{}),
		statsCommandC:             make(chan statsRequest),
		notifyErrorCommandC:       make(chan notifyErrorCommand),
		addrsFromTrackers:         make(chan []*net.TCPAddr),
		addrList:                  addrlist.New(),
		peerIDs:                   make(map[[20]byte]struct{}),
		incomingConnC:             make(chan net.Conn),
		sKeyHash:                  mse.HashSKey(spec.infoHash[:]),
		infoDownloaderResultC:     make(chan infodownloader.Result),
		pieceDownloaderResultC:    make(chan piecedownloader.Result),
		incomingHandshakers:       make(map[string]*incominghandshaker.IncomingHandshaker),
		outgoingHandshakers:       make(map[string]*outgoinghandshaker.OutgoingHandshaker),
		incomingHandshakerResultC: make(chan incominghandshaker.Result),
		outgoingHandshakerResultC: make(chan outgoinghandshaker.Result),
		announcerRequestC:         make(chan *announcer.Request),
		allocatorProgressC:        make(chan allocator.Progress),
		allocatorResultC:          make(chan allocator.Result),
		verifierProgressC:         make(chan verifier.Progress),
		verifierResultC:           make(chan verifier.Result),
	}
	copy(d.peerID[:], peerIDPrefix)
	_, err := rand.Read(d.peerID[len(peerIDPrefix):]) // nolint: gosec
	if err != nil {
		return nil, err
	}
	d.trackersInstances, err = parseTrackers(d.trackers, d.log)
	if err != nil {
		return nil, err
	}
	go d.run()
	return d, nil
}

func parseTrackers(trackers []string, log logger.Logger) ([]tracker.Tracker, error) {
	var ret []tracker.Tracker
	for _, s := range trackers {
		u, err := url.Parse(s)
		if err != nil {
			log.Warningln("cannot parse tracker url:", err)
			continue
		}
		switch u.Scheme {
		case "http", "https":
			ret = append(ret, httptracker.New(u))
		case "udp":
			ret = append(ret, udptracker.New(u))
		default:
			log.Warningln("unsupported tracker scheme: %s", u.Scheme)
		}
	}
	if len(ret) == 0 {
		return nil, errors.New("no tracker found")
	}
	return ret, nil
}
<commit_msg>increase max parallel piece downloads<commit_after>// Package torrent provides a torrent client implementation.
package torrent

import (
	"encoding/hex"
	"errors"
	"io"
	"math/rand"
	"net"
	"net/url"
	"time"

	"github.com/cenkalti/rain/internal/clientversion"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/resume"
	"github.com/cenkalti/rain/storage"
	"github.com/cenkalti/rain/storage/filestorage"
	"github.com/cenkalti/rain/torrent/internal/acceptor"
	"github.com/cenkalti/rain/torrent/internal/addrlist"
	"github.com/cenkalti/rain/torrent/internal/allocator"
	"github.com/cenkalti/rain/torrent/internal/announcer"
	"github.com/cenkalti/rain/torrent/internal/bitfield"
	"github.com/cenkalti/rain/torrent/internal/handshaker/incominghandshaker"
	"github.com/cenkalti/rain/torrent/internal/handshaker/outgoinghandshaker"
	"github.com/cenkalti/rain/torrent/internal/infodownloader"
	"github.com/cenkalti/rain/torrent/internal/magnet"
	"github.com/cenkalti/rain/torrent/internal/metainfo"
	"github.com/cenkalti/rain/torrent/internal/mse"
	"github.com/cenkalti/rain/torrent/internal/peer"
	"github.com/cenkalti/rain/torrent/internal/piece"
	"github.com/cenkalti/rain/torrent/internal/piecedownloader"
	"github.com/cenkalti/rain/torrent/internal/torrentdata"
	"github.com/cenkalti/rain/torrent/internal/tracker"
	"github.com/cenkalti/rain/torrent/internal/tracker/httptracker"
	"github.com/cenkalti/rain/torrent/internal/tracker/udptracker"
	"github.com/cenkalti/rain/torrent/internal/verifier"
)

const (
	parallelInfoDownloads  = 4
	parallelPieceDownloads = 80
	maxPeerDial            = 40
	maxPeerAccept          = 40
)

var (
	// http://www.bittorrent.org/beps/bep_0020.html
	peerIDPrefix = []byte("-RN" + clientversion.Version + "-")
)

// Torrent connects to peers and downloads files from swarm.
type Torrent struct {
	*downloadSpec

	// Unique peer ID is generated per downloader.
	peerID [20]byte

	// Data provides IO access to pieces in torrent.
	data *torrentdata.Data

	// Contains state about the pieces in torrent.
	pieces []*piece.Piece

	// Contains pieces in sorted order for piece selection function.
	sortedPieces []*piece.Piece

	// Peers are sent to this channel when they are disconnected.
	peerDisconnectedC chan *peer.Peer

	// All messages coming from peers are sent to this channel.
	messages chan peer.Message

	// We keep connected peers in this map after they complete handshake phase.
	peers         map[*peer.Peer]struct{}
	incomingPeers map[*peer.Peer]struct{}
	outgoingPeers map[*peer.Peer]struct{}

	// Active piece downloads are kept in this map.
	pieceDownloaders map[*peer.Peer]*piecedownloader.PieceDownloader

	// Active metadata downloads are kept in this map.
	infoDownloaders map[*peer.Peer]*infodownloader.InfoDownloader

	// A peer is optimistically unchoked regardless of their download rate.
	optimisticUnchokedPeer *peer.Peer

	// This channel is closed once all pieces are downloaded and verified.
	completeC chan struct{}

	// If any unrecoverable error occurs, it will be sent to this channel and download will be stopped.
	errC chan error

	// When Stop() is called, it will close this channel to signal run() function to stop.
	closeC chan struct{}

	// This channel will be closed after run loop exists.
	doneC chan struct{}

	// These are the channels for sending a message to run() loop.
	statsCommandC       chan statsRequest       // Stats()
	startCommandC       chan struct{}           // Start()
	stopCommandC        chan struct{}           // Stop()
	notifyErrorCommandC chan notifyErrorCommand // NotifyError()

	// Trackers send announce responses to this channel.
	addrsFromTrackers chan []*net.TCPAddr

	// Keeps a list of peer addresses to connect.
	addrList *addrlist.AddrList

	// New raw connections created by OutgoingHandshaker are sent to here.
	incomingConnC chan net.Conn

	// Keep a set of peer IDs to block duplicate connections.
	peerIDs map[[20]byte]struct{}

	// Listens for incoming peer connections.
	acceptor *acceptor.Acceptor

	// Special hash of info hash for encypted connection handshake.
	sKeyHash [20]byte

	// Tracker implementations for giving to announcers.
	trackersInstances []tracker.Tracker

	// Announces the status of torrent to trackers to get peer addresses.
	announcers []*announcer.Announcer

	// List of peers in handshake state.
	incomingHandshakers map[string]*incominghandshaker.IncomingHandshaker
	outgoingHandshakers map[string]*outgoinghandshaker.OutgoingHandshaker

	// Handshake results are sent to these channels by handshakers.
	incomingHandshakerResultC chan incominghandshaker.Result
	outgoingHandshakerResultC chan outgoinghandshaker.Result

	// When metadata of the torrent downloaded completely, a message is sent to this channel.
	infoDownloaderResultC chan infodownloader.Result

	// When a piece is downloaded completely a message is sent to this channel.
	pieceDownloaderResultC chan piecedownloader.Result

	// Announcers send a request to this channel to get information about the torrent.
	announcerRequestC chan *announcer.Request

	// A timer that ticks periodically to keep a certain number of peers unchoked.
	unchokeTimer  *time.Ticker
	unchokeTimerC <-chan time.Time

	// A timer that ticks periodically to keep a random peer unchoked regardless of its upload rate.
	optimisticUnchokeTimer  *time.Ticker
	optimisticUnchokeTimerC <-chan time.Time

	allocator          *allocator.Allocator
	allocatorProgressC chan allocator.Progress
	allocatorResultC   chan allocator.Result

	verifier          *verifier.Verifier
	verifierProgressC chan verifier.Progress
	verifierResultC   chan verifier.Result

	log logger.Logger
}

// New returns a new torrent by reading a torrent metainfo file.
func New(r io.Reader, port int, sto storage.Storage) (*Torrent, error) {
	m, err := metainfo.New(r)
	if err != nil {
		return nil, err
	}
	spec := &downloadSpec{
		infoHash: m.Info.Hash,
		trackers: m.GetTrackers(),
		name:     m.Info.Name,
		port:     port,
		storage:  sto,
		info:     m.Info,
	}
	return newTorrent(spec)
}

// NewMagnet returns a new torrent by parsing a magnet link.
func NewMagnet(link string, port int, sto storage.Storage) (*Torrent, error) {
	m, err := magnet.New(link)
	if err != nil {
		return nil, err
	}
	spec := &downloadSpec{
		infoHash: m.InfoHash,
		trackers: m.Trackers,
		name:     m.Name,
		port:     port,
		storage:  sto,
	}
	return newTorrent(spec)
}

// NewResume returns a new torrent by loading all info from a resume.DB.
func NewResume(res resume.DB) (*Torrent, error) {
	spec, err := res.Read()
	if err != nil {
		return nil, err
	}
	if spec == nil {
		return nil, errors.New("no resume info")
	}
	return loadResumeSpec(spec, res)
}

// Name of the torrent.
// For magnet downloads name can change after metadata is downloaded but this method still returns the initial name.
func (t *Torrent) Name() string {
	return t.name
}

// InfoHash string encoded in hex.
// InfoHash is a unique value that identifies the files in torrent.
func (t *Torrent) InfoHash() string {
	return hex.EncodeToString(t.infoHash[:])
}

// SetResume adds resume capability to the torrent.
// It must be called before Start() is called.
func (t *Torrent) SetResume(res resume.DB) error {
	spec, err := res.Read()
	if err != nil {
		return err
	}
	if spec == nil {
		return t.writeResume(res)
	}
	t2, err := loadResumeSpec(spec, res)
	if err != nil {
		return err
	}
	if t.InfoHash() != t2.InfoHash() {
		t2.Close()
		return errors.New("invalid resume file (info hashes does not match)")
	}
	t.Close()
	*t = *t2
	return nil
}

func loadResumeSpec(spec *resume.Spec, res resume.DB) (*Torrent, error) {
	var err error
	dspec := &downloadSpec{
		port:     spec.Port,
		trackers: spec.Trackers,
		name:     spec.Name,
		resume:   res,
	}
	copy(dspec.infoHash[:], spec.InfoHash)
	if len(spec.Info) > 0 {
		dspec.info, err = metainfo.NewInfo(spec.Info)
		if err != nil {
			return nil, err
		}
		if len(spec.Bitfield) > 0 {
			dspec.bitfield = bitfield.New(dspec.info.NumPieces)
			copy(dspec.bitfield.Bytes(), spec.Bitfield)
		}
	}
	switch spec.StorageType {
	case filestorage.StorageType:
		dspec.storage = &filestorage.FileStorage{}
	default:
		return nil, errors.New("unknown storage type: " + spec.StorageType)
	}
	err = dspec.storage.Load(spec.StorageArgs)
	if err != nil {
		return nil, err
	}
	return newTorrent(dspec)
}

func (t *Torrent) writeResume(res resume.DB) error {
	rspec := &resume.Spec{
		InfoHash:    t.infoHash[:],
		Port:        t.port,
		Name:        t.name,
		Trackers:    t.trackers,
		StorageType: t.storage.Type(),
		StorageArgs: t.storage.Args(),
	}
	if t.info != nil {
		rspec.Info = t.info.Bytes
	}
	if t.bitfield != nil {
		rspec.Bitfield = t.bitfield.Bytes()
	}
	return res.Write(rspec)
}

func newTorrent(spec *downloadSpec) (*Torrent, error) {
	logName := spec.name
	if len(logName) > 8 {
		logName = logName[:8]
	}
	d := &Torrent{
		downloadSpec:              spec,
		log:                       logger.New("download " + logName),
		peerDisconnectedC:         make(chan *peer.Peer),
		messages:                  make(chan peer.Message),
		peers:                     make(map[*peer.Peer]struct{}),
		incomingPeers:             make(map[*peer.Peer]struct{}),
		outgoingPeers:             make(map[*peer.Peer]struct{}),
		pieceDownloaders:          make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		infoDownloaders:           make(map[*peer.Peer]*infodownloader.InfoDownloader),
		completeC:                 make(chan struct{}),
		closeC:                    make(chan struct{}),
		doneC:                     make(chan struct{}),
		startCommandC:             make(chan struct{}),
		stopCommandC:              make(chan struct{}),
		statsCommandC:             make(chan statsRequest),
		notifyErrorCommandC:       make(chan notifyErrorCommand),
		addrsFromTrackers:         make(chan []*net.TCPAddr),
		addrList:                  addrlist.New(),
		peerIDs:                   make(map[[20]byte]struct{}),
		incomingConnC:             make(chan net.Conn),
		sKeyHash:                  mse.HashSKey(spec.infoHash[:]),
		infoDownloaderResultC:     make(chan infodownloader.Result),
		pieceDownloaderResultC:    make(chan piecedownloader.Result),
		incomingHandshakers:       make(map[string]*incominghandshaker.IncomingHandshaker),
		outgoingHandshakers:       make(map[string]*outgoinghandshaker.OutgoingHandshaker),
		incomingHandshakerResultC: make(chan incominghandshaker.Result),
		outgoingHandshakerResultC: make(chan outgoinghandshaker.Result),
		announcerRequestC:         make(chan *announcer.Request),
		allocatorProgressC:        make(chan allocator.Progress),
		allocatorResultC:          make(chan allocator.Result),
		verifierProgressC:         make(chan verifier.Progress),
		verifierResultC:           make(chan verifier.Result),
	}
	copy(d.peerID[:], peerIDPrefix)
	_, err := rand.Read(d.peerID[len(peerIDPrefix):]) // nolint: gosec
	if err != nil {
		return nil, err
	}
	d.trackersInstances, err = parseTrackers(d.trackers, d.log)
	if err != nil {
		return nil, err
	}
	go d.run()
	return d, nil
}

func parseTrackers(trackers []string, log logger.Logger) ([]tracker.Tracker, error) {
	var ret []tracker.Tracker
	for _, s := range trackers {
		u, err := url.Parse(s)
		if err != nil {
			log.Warningln("cannot parse tracker url:", err)
			continue
		}
		switch u.Scheme {
		case "http", "https":
			ret = append(ret, httptracker.New(u))
		case "udp":
			ret = append(ret, udptracker.New(u))
		default:
			log.Warningln("unsupported tracker scheme: %s", u.Scheme)
		}
	}
	if len(ret) == 0 {
		return nil, errors.New("no tracker found")
	}
	return ret, nil
}
<|endoftext|>