<commit_before>package httpRelay

import (
	"bufio"
	"io"
	"net/http"

	"golang.org/x/net/proxy"
)

// mnot's blog: https://www.mnot.net/blog/2011/07/11/what_proxies_must_do
// rfc: http://tools.ietf.org/html/draft-ietf-httpbis-p1-messaging-14#section-3.3

// HTTPProxyHandler is a proxy handler that passes on request to a SOCKS5 proxy server.
type HTTPProxyHandler struct {
	// Dialer is the dialer for connecting to the SOCKS5 proxy.
	Dialer proxy.Dialer
}

func (h *HTTPProxyHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
	defer func() {
		logError(req.Body.Close(), "error closing client request body:")
	}()
	var err error
	switch req.Method {
	case "CONNECT":
		err = h.handleConnect(resp, req)
	default:
		err = h.processRequest(resp, req)
	}
	logError(err, "Error serving proxy relay")
}

func (h *HTTPProxyHandler) processRequest(resp http.ResponseWriter, req *http.Request) error {
	var err error
	logRequest(req)
	// Verification of requests is already handled by net/http library.
	// Establish connection with socks proxy
	conn, err := h.Dialer.Dial("tcp", fullHost(req.Host))
	if err != nil {
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	defer func() {
		logError(conn.Close(), "error closing connection to socks proxy:")
	}()
	// Prepare request for socks proxy
	proxyReq, err := http.NewRequest(req.Method, req.RequestURI, req.Body)
	if err != nil {
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	// Transfer headers to proxy request
	copyHeaders(proxyReq.Header, req.Header)
	// FIXME add Via header
	// FIXME add what user agent? (Does setting header actually work?)
	proxyReq.Header.Add("User-Agent", "proxy")
	// Send request to socks proxy
	if err = proxyReq.Write(conn); err != nil {
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	// Read proxy response
	proxyRespReader := bufio.NewReader(conn)
	proxyResp, err := http.ReadResponse(proxyRespReader, proxyReq)
	if err != nil {
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	// Transfer headers to client response
	copyHeaders(resp.Header(), proxyResp.Header)
	// Verification of response is already handled by net/http library.
	resp.WriteHeader(proxyResp.StatusCode)
	_, err = io.Copy(resp, proxyResp.Body)
	if err != nil {
		return err
	}
	logError(proxyResp.Body.Close(), "error closing response body:")
	return nil
}

func (h *HTTPProxyHandler) handleConnect(resp http.ResponseWriter, req *http.Request) error {
	logRequest(req)
	// Establish connection with socks proxy
	proxyConn, err := h.Dialer.Dial("tcp", req.Host)
	if err != nil {
		return err
	}
	// Acquire raw connection to the client
	clientConn, err := acquireConn(resp)
	if err != nil {
		logError(proxyConn.Close(), "error while closing proxy connection:")
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	// Send 200 Connection established to client to signal tunnel ready
	_, err = clientConn.Write([]byte("HTTP/1.1 200 Connection established\r\n\r\n\r\n"))
	if err != nil {
		logError(proxyConn.Close(), "error while closing proxy connection:")
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	// Start copying data from one connection to the other
	go transfer(proxyConn, clientConn)
	go transfer(clientConn, proxyConn)
	return nil
}
<commit_msg>Add UserAgent as field in HTTPProxyHandler struct.<commit_after>package httpRelay

import (
	"bufio"
	"io"
	"net/http"

	"golang.org/x/net/proxy"
)

// mnot's blog: https://www.mnot.net/blog/2011/07/11/what_proxies_must_do
// rfc: http://tools.ietf.org/html/draft-ietf-httpbis-p1-messaging-14#section-3.3

// HTTPProxyHandler is a proxy handler that passes on request to a SOCKS5 proxy server.
type HTTPProxyHandler struct {
	// Dialer is the dialer for connecting to the SOCKS5 proxy.
	Dialer    proxy.Dialer
	UserAgent string
}

func (h *HTTPProxyHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
	defer func() {
		logError(req.Body.Close(), "error closing client request body:")
	}()
	var err error
	switch req.Method {
	case "CONNECT":
		err = h.handleConnect(resp, req)
	default:
		err = h.processRequest(resp, req)
	}
	logError(err, "Error serving proxy relay")
}

func (h *HTTPProxyHandler) processRequest(resp http.ResponseWriter, req *http.Request) error {
	var err error
	logRequest(req)
	// Verification of requests is already handled by net/http library.
	// Establish connection with socks proxy
	conn, err := h.Dialer.Dial("tcp", fullHost(req.Host))
	if err != nil {
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	defer func() {
		logError(conn.Close(), "error closing connection to socks proxy:")
	}()
	// Prepare request for socks proxy
	proxyReq, err := http.NewRequest(req.Method, req.RequestURI, req.Body)
	if err != nil {
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	// Transfer headers to proxy request
	copyHeaders(proxyReq.Header, req.Header)
	// FIXME add Via header
	if h.UserAgent != "" {
		// Add specified user agent as header.
		proxyReq.Header.Add("User-Agent", h.UserAgent)
	}
	// Send request to socks proxy
	if err = proxyReq.Write(conn); err != nil {
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	// Read proxy response
	proxyRespReader := bufio.NewReader(conn)
	proxyResp, err := http.ReadResponse(proxyRespReader, proxyReq)
	if err != nil {
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	// Transfer headers to client response
	copyHeaders(resp.Header(), proxyResp.Header)
	// Verification of response is already handled by net/http library.
	resp.WriteHeader(proxyResp.StatusCode)
	_, err = io.Copy(resp, proxyResp.Body)
	if err != nil {
		return err
	}
	logError(proxyResp.Body.Close(), "error closing response body:")
	return nil
}

func (h *HTTPProxyHandler) handleConnect(resp http.ResponseWriter, req *http.Request) error {
	logRequest(req)
	// Establish connection with socks proxy
	proxyConn, err := h.Dialer.Dial("tcp", req.Host)
	if err != nil {
		return err
	}
	// Acquire raw connection to the client
	clientConn, err := acquireConn(resp)
	if err != nil {
		logError(proxyConn.Close(), "error while closing proxy connection:")
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	// Send 200 Connection established to client to signal tunnel ready
	_, err = clientConn.Write([]byte("HTTP/1.1 200 Connection established\r\n\r\n\r\n"))
	if err != nil {
		logError(proxyConn.Close(), "error while closing proxy connection:")
		resp.WriteHeader(http.StatusInternalServerError)
		return err
	}
	// Start copying data from one connection to the other
	go transfer(proxyConn, clientConn)
	go transfer(clientConn, proxyConn)
	return nil
}
<|endoftext|>


<commit_before>package components

import (
	"github.com/flosch/pongo2"
	"github.com/golang/glog"

	"bitbucket.org/moovie/util/stringslice"
	"bitbucket.org/moovie/util/template"
)

// Render - Renders compiled component.
// Only first template context is accepted.
// Sets source component in template context under key `source_component`.
func Render(c *Compiled, ctxs ...template.Context) (res *Rendered, err error) {
	var ctx template.Context
	if len(ctxs) == 0 || ctxs[0] == nil {
		ctx = make(template.Context)
	} else {
		ctx = ctxs[0]
	}
	ctx["source_component"] = c.Component
	res = new(Rendered)
	err = renderComponent(c, res, res, ctx)
	return
}

// renderComponent - Renders a component.
// `main` is where `Styles` and `Scripts` are inserted.
// `res` is where `Body` is inserted.
func renderComponent(c *Compiled, main, res *Rendered, ctx template.Context) (err error) {
	// Set component defaults
	ctx, err = withComponentDefaults(c, ctx)
	if err != nil {
		return
	}

	if glog.V(6) {
		glog.Infof("[render] name=%q ctx=%#v", c.Name, ctx)
	}

	// Render required components
	for name, req := range c.Require {
		r := new(Rendered)
		err = renderComponent(req, main, r, ctx)
		if err != nil {
			return
		}

		// Add resulting body to context
		ctx[name] = pongo2.AsSafeValue(r.Body)
	}

	// Render `Main` component template
	res.Body, err = template.ExecuteToString(c.Main, ctx)
	if err != nil {
		return
	}

	// Extend a template if any
	if c.Extends != nil {
		ctx["children"] = pongo2.AsSafeValue(res.Body)
		err = renderComponent(c.Extends, main, res, ctx)
		if err != nil {
			return
		}
	}

	// Render component styles and scripts
	err = renderAssets(c, main, ctx)
	if err != nil {
		return
	}

	return
}

func renderAssets(c *Compiled, res *Rendered, ctx template.Context) (err error) {
	// Render component styles
	tmp, err := template.ExecuteList(c.Styles, ctx)
	if err != nil {
		return
	}

	// Merge component scripts into result
	res.Styles = stringslice.MergeUnique(res.Styles, tmp)

	// Render component scripts
	tmp, err = template.ExecuteList(c.Scripts, ctx)
	if err != nil {
		return
	}

	// Merge component scripts into result
	res.Scripts = stringslice.MergeUnique(res.Scripts, tmp)
	return
}

// withComponentDefaults - Returns a context with component defaults set.
func withComponentDefaults(c *Compiled, ctx template.Context) (_ template.Context, err error) {
	// Set defaults from component base context
	ctx = ctx.WithDefaults(c.Context)

	// Return if no `With` to merge with
	if len(c.With) == 0 {
		return ctx, nil
	}

	// Execute component's `With` templates
	w, err := c.With.Execute(ctx)
	if err != nil {
		return
	}

	// Merge compiled `With` into context
	ctx = ctx.Merge(w)
	return ctx, nil
}
<commit_msg>fix merging  to Context as defaults<commit_after>package components

import (
	"github.com/flosch/pongo2"
	"github.com/golang/glog"

	"bitbucket.org/moovie/util/stringslice"
	"bitbucket.org/moovie/util/template"
)

// Render - Renders compiled component.
// Only first template context is accepted.
// Sets source component in template context under key `source_component`.
func Render(c *Compiled, ctxs ...template.Context) (res *Rendered, err error) {
	var ctx template.Context
	if len(ctxs) == 0 || ctxs[0] == nil {
		ctx = make(template.Context)
	} else {
		ctx = ctxs[0]
	}
	ctx["source_component"] = c.Component
	res = new(Rendered)
	err = renderComponent(c, res, res, ctx)
	return
}

// renderComponent - Renders a component.
// `main` is where `Styles` and `Scripts` are inserted.
// `res` is where `Body` is inserted.
func renderComponent(c *Compiled, main, res *Rendered, ctx template.Context) (err error) {
	// Set component defaults
	ctx, err = withComponentDefaults(c, ctx)
	if err != nil {
		return
	}

	if glog.V(6) {
		glog.Infof("[render] name=%q ctx=%#v", c.Name, ctx)
	}

	// Render required components
	for name, req := range c.Require {
		r := new(Rendered)
		err = renderComponent(req, main, r, ctx)
		if err != nil {
			return
		}

		// Add resulting body to context
		ctx[name] = pongo2.AsSafeValue(r.Body)
	}

	// Render `Main` component template
	res.Body, err = template.ExecuteToString(c.Main, ctx)
	if err != nil {
		return
	}

	// Extend a template if any
	if c.Extends != nil {
		ctx["children"] = pongo2.AsSafeValue(res.Body)
		err = renderComponent(c.Extends, main, res, ctx)
		if err != nil {
			return
		}
	}

	// Render component styles and scripts
	err = renderAssets(c, main, ctx)
	if err != nil {
		return
	}

	return
}

func renderAssets(c *Compiled, res *Rendered, ctx template.Context) (err error) {
	// Render component styles
	tmp, err := template.ExecuteList(c.Styles, ctx)
	if err != nil {
		return
	}

	// Merge component scripts into result
	res.Styles = stringslice.MergeUnique(res.Styles, tmp)

	// Render component scripts
	tmp, err = template.ExecuteList(c.Scripts, ctx)
	if err != nil {
		return
	}

	// Merge component scripts into result
	res.Scripts = stringslice.MergeUnique(res.Scripts, tmp)
	return
}

// withComponentDefaults - Returns a context with component defaults set.
func withComponentDefaults(c *Compiled, ctx template.Context) (_ template.Context, err error) {
	// Set defaults from component base context
	ctx = ctx.WithDefaults(c.Context)

	// Return if no `With` to merge with
	if len(c.With) == 0 {
		return ctx, nil
	}

	// Execute component's `With` templates
	for key, node := range c.With {
		_, has := ctx[key]
		if has {
			continue
		}
		ctx[key], err = node.Execute(ctx)
		if err != nil {
			return
		}
	}

	return ctx, nil
}
<|endoftext|>


<commit_before>package gardenhealth

import (
	"os"
	"time"

	"github.com/cloudfoundry-incubator/executor"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
)

type HealthcheckTimeoutError struct{}

func (HealthcheckTimeoutError) Error() string {
	return "garden healthcheck timed out"
}

//go:generate counterfeiter -o fakegardenhealth/fake_timerprovider.go . TimerProvider

type TimerProvider interface {
	NewTimer(time.Duration) clock.Timer
}

type Runner struct {
	failures        int
	healthy         bool
	checkInterval   time.Duration
	timeoutInterval time.Duration
	logger          lager.Logger
	checker         Checker
	executorClient  executor.Client
	timerProvider   TimerProvider
}

func NewRunner(
	checkInterval time.Duration,
	timeoutInterval time.Duration,
	logger lager.Logger,
	checker Checker,
	executorClient executor.Client,
	timerProvider TimerProvider,
) *Runner {
	return &Runner{
		checkInterval:   checkInterval,
		timeoutInterval: timeoutInterval,
		logger:          logger.Session("garden-healthcheck"),
		checker:         checker,
		executorClient:  executorClient,
		timerProvider:   timerProvider,
		healthy:         false,
		failures:        0,
	}
}

func (r *Runner) Run(signals <-chan os.Signal, ready chan<- struct{}) error {
	var (
		startHealthcheck    = r.timerProvider.NewTimer(0)
		healthcheckTimeout  = r.timerProvider.NewTimer(r.timeoutInterval)
		healthcheckComplete = make(chan error, 1)
	)
	r.logger.Info("starting")

	go r.HealthcheckCycle(healthcheckComplete)

	select {
	case <-signals:
		return nil

	case <-healthcheckTimeout.C():
		r.logger.Error("failed-initial-healthcheck-timeout", nil)
		return HealthcheckTimeoutError{}

	case err := <-healthcheckComplete:
		if err != nil {
			r.logger.Error("failed-initial-healthcheck", err)
			return err
		}
		healthcheckTimeout.Stop()
	}

	r.logger.Info("passed-initial-healthcheck")
	r.SetHealthy()

	close(ready)
	r.logger.Info("started")

	startHealthcheck.Reset(r.checkInterval)

	for {
		select {
		case <-signals:
			r.logger.Info("complete")
			return nil

		case <-startHealthcheck.C():
			r.logger.Info("check-starting")
			go r.HealthcheckCycle(healthcheckComplete)
			healthcheckTimeout.Reset(r.timeoutInterval)

		case <-healthcheckTimeout.C():
			r.logger.Error("failed-healthcheck-timeout", nil)
			r.SetUnhealthy()

		case err := <-healthcheckComplete:
			timeoutOk := healthcheckTimeout.Stop()
			switch err.(type) {
			case nil:
				r.logger.Info("passed-health-check")
				if timeoutOk {
					r.SetHealthy()
				}

			default:
				r.logger.Error("failed-health-check", err)
				r.SetUnhealthy()
			}

			startHealthcheck.Reset(r.checkInterval)
			r.logger.Info("check-complete")
		}
	}
}

func (r *Runner) SetHealthy() {
	if !r.healthy {
		r.logger.Info("set-state-healthy")
		r.executorClient.SetHealthy(true)
		r.healthy = true
	}
}

func (r *Runner) SetUnhealthy() {
	if r.healthy {
		r.logger.Error("set-state-unhealthy", nil)
		r.executorClient.SetHealthy(false)
		r.healthy = false
	}
}

func (r *Runner) HealthcheckCycle(healthcheckComplete chan<- error) {
	healthcheckComplete <- r.checker.Healthcheck(r.logger)
}
<commit_msg>Don't export private methods<commit_after>package gardenhealth

import (
	"os"
	"time"

	"github.com/cloudfoundry-incubator/executor"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
)

type HealthcheckTimeoutError struct{}

func (HealthcheckTimeoutError) Error() string {
	return "garden healthcheck timed out"
}

//go:generate counterfeiter -o fakegardenhealth/fake_timerprovider.go . TimerProvider

type TimerProvider interface {
	NewTimer(time.Duration) clock.Timer
}

type Runner struct {
	failures        int
	healthy         bool
	checkInterval   time.Duration
	timeoutInterval time.Duration
	logger          lager.Logger
	checker         Checker
	executorClient  executor.Client
	timerProvider   TimerProvider
}

func NewRunner(
	checkInterval time.Duration,
	timeoutInterval time.Duration,
	logger lager.Logger,
	checker Checker,
	executorClient executor.Client,
	timerProvider TimerProvider,
) *Runner {
	return &Runner{
		checkInterval:   checkInterval,
		timeoutInterval: timeoutInterval,
		logger:          logger.Session("garden-healthcheck"),
		checker:         checker,
		executorClient:  executorClient,
		timerProvider:   timerProvider,
		healthy:         false,
		failures:        0,
	}
}

func (r *Runner) Run(signals <-chan os.Signal, ready chan<- struct{}) error {
	var (
		startHealthcheck    = r.timerProvider.NewTimer(0)
		healthcheckTimeout  = r.timerProvider.NewTimer(r.timeoutInterval)
		healthcheckComplete = make(chan error, 1)
	)
	r.logger.Info("starting")

	go r.healthcheckCycle(healthcheckComplete)

	select {
	case <-signals:
		return nil

	case <-healthcheckTimeout.C():
		r.logger.Error("failed-initial-healthcheck-timeout", nil)
		return HealthcheckTimeoutError{}

	case err := <-healthcheckComplete:
		if err != nil {
			r.logger.Error("failed-initial-healthcheck", err)
			return err
		}
		healthcheckTimeout.Stop()
	}

	r.logger.Info("passed-initial-healthcheck")
	r.setHealthy()

	close(ready)
	r.logger.Info("started")

	startHealthcheck.Reset(r.checkInterval)

	for {
		select {
		case <-signals:
			r.logger.Info("complete")
			return nil

		case <-startHealthcheck.C():
			r.logger.Info("check-starting")
			go r.healthcheckCycle(healthcheckComplete)
			healthcheckTimeout.Reset(r.timeoutInterval)

		case <-healthcheckTimeout.C():
			r.logger.Error("failed-healthcheck-timeout", nil)
			r.setUnhealthy()

		case err := <-healthcheckComplete:
			timeoutOk := healthcheckTimeout.Stop()
			switch err.(type) {
			case nil:
				r.logger.Info("passed-health-check")
				if timeoutOk {
					r.setHealthy()
				}

			default:
				r.logger.Error("failed-health-check", err)
				r.setUnhealthy()
			}

			startHealthcheck.Reset(r.checkInterval)
			r.logger.Info("check-complete")
		}
	}
}

func (r *Runner) setHealthy() {
	if !r.healthy {
		r.logger.Info("set-state-healthy")
		r.executorClient.SetHealthy(true)
		r.healthy = true
	}
}

func (r *Runner) setUnhealthy() {
	if r.healthy {
		r.logger.Error("set-state-unhealthy", nil)
		r.executorClient.SetHealthy(false)
		r.healthy = false
	}
}

func (r *Runner) healthcheckCycle(healthcheckComplete chan<- error) {
	healthcheckComplete <- r.checker.Healthcheck(r.logger)
}
<|endoftext|>


<commit_before>package backend

import (
	"fmt"

	"github.com/bitterfly/kuho/spiderdata"
	"github.com/jmoiron/sqlx"
	_ "github.com/lib/pq"
)

type Backend struct {
	db *sqlx.DB
}

func New(dbURN string) (*Backend, error) {
	var db *sqlx.DB
	db, err := sqlx.Connect("postgres", dbURN)
	if err != nil {
		return nil, fmt.Errorf("unable to connect to database: %s", err)
	}

	return &Backend{
		db: db,
	}, nil
}

func (b *Backend) Foo() (string, error) {
	return "this is foo.", nil
}

func (b *Backend) InitDB() error {
	_, err := b.db.Exec(schema)
	if err != nil {
		return fmt.Errorf("unable to execute schema: %s", err)
	}
	return nil
}

func (b *Backend) DropDB() error {
	_, err := b.db.Exec(dropSchema)
	if err != nil {
		return fmt.Errorf("unable to drop schema: %s", err)
	}
	return nil
}

func (b *Backend) Fill(data *spiderdata.Request) error {
	_, err := b.Wrap(func(tx *sqlx.Tx) (interface{}, error) {
		// first clear all of the old data

		cinemas := data.Cinemas
		films := data.Films

		// insert cinemas
		for i := range cinemas {
			err := insertCinema(tx, cinemas[i])
			if err != nil {
				return nil, fmt.Errorf("unable to insert cinema: %s", err)
			}
		}

		// insert films
		for i := range films {
			err := insertFilm(tx, films[i])
			if err != nil {
				return nil, fmt.Errorf("unable to insert timetable: %s", err)
			}
		}
		return nil, nil
	})
	return err
}

func insertCinema(tx *sqlx.Tx, cinema *spiderdata.Cinema) error {
	//name, url, chain, lastUpdate
	_, err := tx.Exec(INSERT_INTO_CINEMA, cinema.Name, cinema.ShortName, cinema.URL, cinema.Chain, &cinema.Acquired)

	if err != nil {
		return fmt.Errorf("Could not write into cinema - %s", err)
	}
	return nil
}

func insertFilm(tx *sqlx.Tx, film *spiderdata.Film) error {
	screenings := film.Screenings

	var filmID int64
	var err error
	//imdbFilmId, title, year,rating, imdbCertainty
	if film.ImdbID == nil {
		err = tx.Get(&filmID, INSERT_INTO_FILM_NULL_IMDBID, nil, film.Title, film.Year, film.Rating, film.ImdbIDCertainty)
	} else {
		err = tx.Get(&filmID, INSERT_INTO_FILM_NOT_NULL_IMDBID, nil, film.Title, film.Year, film.Rating, film.ImdbIDCertainty)
	}
	if err != nil {
		return fmt.Errorf("Could not write into film - %s", err)
	}

	var screeningID int64
	var cinemaId int64
	for i := range screenings {
		screening := screenings[i]
		cinemaShortName := screening.CinemaShortName

		err = tx.Get(&cinemaId, GET_CINEMA_BY_SHORT_NAME, cinemaShortName)
		if err != nil {
			return fmt.Errorf("Could not write get id for cinema %s  - %s", cinemaShortName, err)
		}

		// cinemaId, filmId, hall, duration, language,
		//					  completeTitle, isActive, 	hasSubtitles, hasDub, isImax, is3D, is4D, lastUpdate

		err = tx.Get(&screeningID, INSERT_INTO_SCREENING,
			cinemaId, filmID, screening.Hall, screening.Duration.Nanoseconds(), screening.Language,
			screening.Variant, screening.Active, screening.IsSubtitled, screening.IsDubbed,
			screening.IsImax, screening.Is3D, screening.Is4D, screening.Acquired,
		)

		if err != nil {
			return fmt.Errorf("Could not write into screening - %s", err)
		}

		screeningTickets := screening.Tickets

		for j := range screeningTickets {
			ticket := screeningTickets[j]

			//screeningId, cinemaId, type, bookingURL, price, currency, lastUpdate
			_, err = tx.Exec(INSERT_INTO_TICKET, screeningID, cinemaId,
				ticket.Type, ticket.BookingURL, ticket.Price, ticket.Currency, ticket.Acquired,
			)

			if err != nil {
				return fmt.Errorf("Could not write into ticket - %s", err)

			}

		}
	}
	return nil
}
<commit_msg>enabled usage of imdbIDs but the now don't work because of a constraint<commit_after>package backend

import (
	"fmt"

	"github.com/bitterfly/kuho/spiderdata"
	"github.com/jmoiron/sqlx"
	_ "github.com/lib/pq"
)

type Backend struct {
	db *sqlx.DB
}

func New(dbURN string) (*Backend, error) {
	var db *sqlx.DB
	db, err := sqlx.Connect("postgres", dbURN)
	if err != nil {
		return nil, fmt.Errorf("unable to connect to database: %s", err)
	}

	return &Backend{
		db: db,
	}, nil
}

func (b *Backend) Foo() (string, error) {
	return "this is foo.", nil
}

func (b *Backend) InitDB() error {
	_, err := b.db.Exec(schema)
	if err != nil {
		return fmt.Errorf("unable to execute schema: %s", err)
	}
	return nil
}

func (b *Backend) DropDB() error {
	_, err := b.db.Exec(dropSchema)
	if err != nil {
		return fmt.Errorf("unable to drop schema: %s", err)
	}
	return nil
}

func (b *Backend) Fill(data *spiderdata.Request) error {
	_, err := b.Wrap(func(tx *sqlx.Tx) (interface{}, error) {
		// first clear all of the old data

		cinemas := data.Cinemas
		films := data.Films

		// insert cinemas
		for i := range cinemas {
			err := insertCinema(tx, cinemas[i])
			if err != nil {
				return nil, fmt.Errorf("unable to insert cinema: %s", err)
			}
		}

		// insert films
		for i := range films {
			err := insertFilm(tx, films[i])
			if err != nil {
				return nil, fmt.Errorf("unable to insert timetable: %s", err)
			}
		}
		return nil, nil
	})
	return err
}

func insertCinema(tx *sqlx.Tx, cinema *spiderdata.Cinema) error {
	//name, url, chain, lastUpdate
	_, err := tx.Exec(INSERT_INTO_CINEMA, cinema.Name, cinema.ShortName, cinema.URL, cinema.Chain, &cinema.Acquired)

	if err != nil {
		return fmt.Errorf("Could not write into cinema - %s", err)
	}
	return nil
}

func insertFilm(tx *sqlx.Tx, film *spiderdata.Film) error {
	screenings := film.Screenings

	var filmID int64
	var err error
	//imdbFilmId, title, year,rating, imdbCertainty
	if film.ImdbID == nil {
		err = tx.Get(&filmID, INSERT_INTO_FILM_NULL_IMDBID, film.ImdbID, film.Title, film.Year, film.Rating, film.ImdbIDCertainty)
	} else {
		err = tx.Get(&filmID, INSERT_INTO_FILM_NOT_NULL_IMDBID, film.ImdbID, film.Title, film.Year, film.Rating, film.ImdbIDCertainty)
	}
	if err != nil {
		return fmt.Errorf("Could not write into film - %s", err)
	}

	var screeningID int64
	var cinemaId int64
	for i := range screenings {
		screening := screenings[i]
		cinemaShortName := screening.CinemaShortName

		err = tx.Get(&cinemaId, GET_CINEMA_BY_SHORT_NAME, cinemaShortName)
		if err != nil {
			return fmt.Errorf("Could not write get id for cinema %s  - %s", cinemaShortName, err)
		}

		// cinemaId, filmId, hall, duration, language,
		//					  completeTitle, isActive, 	hasSubtitles, hasDub, isImax, is3D, is4D, lastUpdate

		err = tx.Get(&screeningID, INSERT_INTO_SCREENING,
			cinemaId, filmID, screening.Hall, screening.Duration.Nanoseconds(), screening.Language,
			screening.Variant, screening.Active, screening.IsSubtitled, screening.IsDubbed,
			screening.IsImax, screening.Is3D, screening.Is4D, screening.Acquired,
		)

		if err != nil {
			return fmt.Errorf("Could not write into screening - %s", err)
		}

		screeningTickets := screening.Tickets

		for j := range screeningTickets {
			ticket := screeningTickets[j]

			//screeningId, cinemaId, type, bookingURL, price, currency, lastUpdate
			_, err = tx.Exec(INSERT_INTO_TICKET, screeningID, cinemaId,
				ticket.Type, ticket.BookingURL, ticket.Price, ticket.Currency, ticket.Acquired,
			)

			if err != nil {
				return fmt.Errorf("Could not write into ticket - %s", err)

			}

		}
	}
	return nil
}
<|endoftext|>


<commit_before>package client

import (
	"bytes"
	"encoding/gob"
	"github.com/Symantec/Dominator/lib/hash"
	"github.com/Symantec/Dominator/lib/objectserver"
	"github.com/Symantec/Dominator/lib/srpc"
	proto "github.com/Symantec/Dominator/proto/filegenerator"
	"io"
	"log"
	"time"
)

func newManager(objSrv objectserver.ObjectServer, logger *log.Logger) *Manager {
	sourceConnectChannel := make(chan string)
	m := &Manager{
		sourceMap:            make(map[string]*sourceType),
		objectServer:         objSrv,
		machineMap:           make(map[string]*machineType),
		addMachineChannel:    make(chan *machineType),
		removeMachineChannel: make(chan string),
		updateMachineChannel: make(chan *machineType),
		serverMessageChannel: make(chan *serverMessageType),
		sourceConnectChannel: sourceConnectChannel,
		objectWaiters:        make(map[hash.Hash][]chan<- hash.Hash),
		logger:               logger}
	go m.manage(sourceConnectChannel)
	return m
}

func (m *Manager) manage(sourceConnectChannel <-chan string) {
	for {
		select {
		case machine := <-m.addMachineChannel:
			m.addMachine(machine)
		case hostname := <-m.removeMachineChannel:
			m.removeMachine(hostname)
		case machine := <-m.updateMachineChannel:
			m.updateMachine(machine)
		case serverMessage := <-m.serverMessageChannel:
			m.processMessage(serverMessage)
		case sourceName := <-sourceConnectChannel:
			m.processSourceConnect(sourceName)
		}
	}
}

func (m *Manager) processMessage(serverMessage *serverMessageType) {
	if msg := serverMessage.serverMessage.GetObjectResponse; msg != nil {
		if _, _, err := m.objectServer.AddObject(
			bytes.NewReader(msg.Data), 0, &msg.Hash); err != nil {
			m.logger.Println(err)
		} else {
			if waiters, ok := m.objectWaiters[msg.Hash]; ok {
				for _, channel := range waiters {
					channel <- msg.Hash
				}
				delete(m.objectWaiters, msg.Hash)
			}
		}
	}
	if msg := serverMessage.serverMessage.YieldResponse; msg != nil {
		if machine, ok := m.machineMap[msg.Hostname]; ok {
			m.handleYieldResponse(machine, msg.Files)
		} // else machine no longer known. Drop the message.
	}
}

func (m *Manager) processSourceConnect(sourceName string) {
	source := m.sourceMap[sourceName]
	for _, machine := range m.machineMap {
		if pathnames, ok := machine.sourceToPaths[sourceName]; ok {
			request := &proto.ClientRequest{
				YieldRequest: &proto.YieldRequest{machine.machine, pathnames}}
			source.sendChannel <- request
		}
	}
}

// Returns true if the source was already set up.
func (m *Manager) getSource(sourceName string) (*sourceType, bool) {
	source, ok := m.sourceMap[sourceName]
	if ok {
		return source, true
	}
	source = new(sourceType)
	sendChannel := make(chan *proto.ClientRequest, 4096)
	source.sendChannel = sendChannel
	m.sourceMap[sourceName] = source
	go manageSource(sourceName, m.sourceConnectChannel, sendChannel,
		m.serverMessageChannel, m.logger)
	return source, false
}

func manageSource(sourceName string, sourceConnectChannel chan<- string,
	clientRequestChannel <-chan *proto.ClientRequest,
	serverMessageChannel chan<- *serverMessageType, logger *log.Logger) {
	closeNotifyChannel := make(chan struct{})
	initialTimeout := time.Second
	timeout := initialTimeout
	for ; ; time.Sleep(timeout) {
		if timeout < time.Minute {
			timeout *= 2
		}
		client, err := srpc.DialHTTP("tcp", sourceName, timeout)
		if err != nil {
			logger.Printf("error connecting to: %s: %s\n", sourceName, err)
			continue
		}
		conn, err := client.Call("FileGenerator.Connect")
		if err != nil {
			client.Close()
			logger.Println(err)
			continue
		}
		timeout = initialTimeout
		// The server keeps the same encoder/decoder pair over the lifetime of
		// the connection, so we must do the same.
		go handleServerMessages(sourceName, gob.NewDecoder(conn),
			serverMessageChannel, closeNotifyChannel, logger)
		sourceConnectChannel <- sourceName
		sendClientRequests(conn, clientRequestChannel, closeNotifyChannel,
			logger)
		conn.Close()
		client.Close()
	}
}

func sendClientRequests(conn *srpc.Conn,
	clientRequestChannel <-chan *proto.ClientRequest,
	closeNotifyChannel <-chan struct{}, logger *log.Logger) {
	encoder := gob.NewEncoder(conn)
	for {
		select {
		case clientRequest := <-clientRequestChannel:
			if err := encoder.Encode(clientRequest); err != nil {
				logger.Printf("error encoding client request: %s\n", err)
				return
			}
			if len(clientRequestChannel) < 1 {
				if err := conn.Flush(); err != nil {
					logger.Printf("error flushing: %s\n", err)
					return
				}
			}
		case <-closeNotifyChannel:
			return
		}
	}
}

func handleServerMessages(sourceName string, decoder *gob.Decoder,
	serverMessageChannel chan<- *serverMessageType,
	closeNotifyChannel chan<- struct{}, logger *log.Logger) {
	for {
		var message proto.ServerMessage
		if err := decoder.Decode(&message); err != nil {
			if err == io.EOF {
				logger.Printf("connection to source: %s closed\n", sourceName)
			} else {
				logger.Println(err)
			}
			closeNotifyChannel <- struct{}{}
			return
		}
		serverMessageChannel <- &serverMessageType{sourceName, message}
	}
}
<commit_msg>Change source retry timeout in lib/filegen/client.<commit_after>package client

import (
	"bytes"
	"encoding/gob"
	"github.com/Symantec/Dominator/lib/hash"
	"github.com/Symantec/Dominator/lib/objectserver"
	"github.com/Symantec/Dominator/lib/srpc"
	proto "github.com/Symantec/Dominator/proto/filegenerator"
	"io"
	"log"
	"time"
)

func newManager(objSrv objectserver.ObjectServer, logger *log.Logger) *Manager {
	sourceConnectChannel := make(chan string)
	m := &Manager{
		sourceMap:            make(map[string]*sourceType),
		objectServer:         objSrv,
		machineMap:           make(map[string]*machineType),
		addMachineChannel:    make(chan *machineType),
		removeMachineChannel: make(chan string),
		updateMachineChannel: make(chan *machineType),
		serverMessageChannel: make(chan *serverMessageType),
		sourceConnectChannel: sourceConnectChannel,
		objectWaiters:        make(map[hash.Hash][]chan<- hash.Hash),
		logger:               logger}
	go m.manage(sourceConnectChannel)
	return m
}

func (m *Manager) manage(sourceConnectChannel <-chan string) {
	for {
		select {
		case machine := <-m.addMachineChannel:
			m.addMachine(machine)
		case hostname := <-m.removeMachineChannel:
			m.removeMachine(hostname)
		case machine := <-m.updateMachineChannel:
			m.updateMachine(machine)
		case serverMessage := <-m.serverMessageChannel:
			m.processMessage(serverMessage)
		case sourceName := <-sourceConnectChannel:
			m.processSourceConnect(sourceName)
		}
	}
}

func (m *Manager) processMessage(serverMessage *serverMessageType) {
	if msg := serverMessage.serverMessage.GetObjectResponse; msg != nil {
		if _, _, err := m.objectServer.AddObject(
			bytes.NewReader(msg.Data), 0, &msg.Hash); err != nil {
			m.logger.Println(err)
		} else {
			if waiters, ok := m.objectWaiters[msg.Hash]; ok {
				for _, channel := range waiters {
					channel <- msg.Hash
				}
				delete(m.objectWaiters, msg.Hash)
			}
		}
	}
	if msg := serverMessage.serverMessage.YieldResponse; msg != nil {
		if machine, ok := m.machineMap[msg.Hostname]; ok {
			m.handleYieldResponse(machine, msg.Files)
		} // else machine no longer known. Drop the message.
	}
}

func (m *Manager) processSourceConnect(sourceName string) {
	source := m.sourceMap[sourceName]
	for _, machine := range m.machineMap {
		if pathnames, ok := machine.sourceToPaths[sourceName]; ok {
			request := &proto.ClientRequest{
				YieldRequest: &proto.YieldRequest{machine.machine, pathnames}}
			source.sendChannel <- request
		}
	}
}

// Returns true if the source was already set up.
func (m *Manager) getSource(sourceName string) (*sourceType, bool) {
	source, ok := m.sourceMap[sourceName]
	if ok {
		return source, true
	}
	source = new(sourceType)
	sendChannel := make(chan *proto.ClientRequest, 4096)
	source.sendChannel = sendChannel
	m.sourceMap[sourceName] = source
	go manageSource(sourceName, m.sourceConnectChannel, sendChannel,
		m.serverMessageChannel, m.logger)
	return source, false
}

func manageSource(sourceName string, sourceConnectChannel chan<- string,
	clientRequestChannel <-chan *proto.ClientRequest,
	serverMessageChannel chan<- *serverMessageType, logger *log.Logger) {
	closeNotifyChannel := make(chan struct{})
	initialRetryTimeout := time.Millisecond * 100
	retryTimeout := initialRetryTimeout
	for ; ; time.Sleep(retryTimeout) {
		if retryTimeout < time.Minute {
			retryTimeout *= 2
		}
		client, err := srpc.DialHTTP("tcp", sourceName, time.Second*15)
		if err != nil {
			logger.Printf("error connecting to: %s: %s\n", sourceName, err)
			continue
		}
		conn, err := client.Call("FileGenerator.Connect")
		if err != nil {
			client.Close()
			logger.Println(err)
			continue
		}
		retryTimeout = initialRetryTimeout
		// The server keeps the same encoder/decoder pair over the lifetime of
		// the connection, so we must do the same.
		go handleServerMessages(sourceName, gob.NewDecoder(conn),
			serverMessageChannel, closeNotifyChannel, logger)
		sourceConnectChannel <- sourceName
		sendClientRequests(conn, clientRequestChannel, closeNotifyChannel,
			logger)
		conn.Close()
		client.Close()
	}
}

func sendClientRequests(conn *srpc.Conn,
	clientRequestChannel <-chan *proto.ClientRequest,
	closeNotifyChannel <-chan struct{}, logger *log.Logger) {
	encoder := gob.NewEncoder(conn)
	for {
		select {
		case clientRequest := <-clientRequestChannel:
			if err := encoder.Encode(clientRequest); err != nil {
				logger.Printf("error encoding client request: %s\n", err)
				return
			}
			if len(clientRequestChannel) < 1 {
				if err := conn.Flush(); err != nil {
					logger.Printf("error flushing: %s\n", err)
					return
				}
			}
		case <-closeNotifyChannel:
			return
		}
	}
}

func handleServerMessages(sourceName string, decoder *gob.Decoder,
	serverMessageChannel chan<- *serverMessageType,
	closeNotifyChannel chan<- struct{}, logger *log.Logger) {
	for {
		var message proto.ServerMessage
		if err := decoder.Decode(&message); err != nil {
			if err == io.EOF {
				logger.Printf("connection to source: %s closed\n", sourceName)
			} else {
				logger.Println(err)
			}
			closeNotifyChannel <- struct{}{}
			return
		}
		serverMessageChannel <- &serverMessageType{sourceName, message}
	}
}
<|endoftext|>


<commit_before>// Package backend implements the persistent AzmoDB backend key/value
// database.
package backend

import (
	"crypto/sha1"
	"errors"
	"io"
	"time"

	btree "github.com/boltdb/bolt"
	"github.com/golang/snappy"
)

// Backend represents a persistent AzmoDB backend.
type Backend interface {
	// Range performs fn on all values stored in the database at rev.
	// If fn returns an errors the traversal is stopped and the error
	// is returned.
	Range(rev Revision, fn func(key []byte, value []byte) error) error

	// Batch starts a new batch transaction. Starting multiple write
	// batch transactions will cause the calls to block and be
	// serialized until the current write batch transaction finishes.
	Batch(rev Revision) (Batch, error)

	// Last returns the last revision in the database and an error if
	// any.
	Last() (Revision, error)
}

// Batch returns a batch transaction on the database.
type Batch interface {
	// Put sets the value for a key in the database. Put must create a
	// copy of the supplied key and value.
	Put(key []byte, value []byte) error

	// Close closes the batch transaction.
	Close() error
}

// Revision represents a serialized AzmoDB revision.
type Revision [8]byte

var (
	rootBuckets = [][]byte{dataBucket, metaBucket}
	dataBucket  = []byte("__data__")
	metaBucket  = []byte("__meta__")

	_ Backend = (*DB)(nil)
)

// Option represents a DB option function.
type Option func(*DB) error

// WithMaxBatchEntries configures the maximum batch entries.
func WithMaxBatchEntries(entries int) Option {
	return func(db *DB) error {
		db.maxEntries = entries
		return nil
	}
}

// WithMaxBatchSize configures the maximum batch size.
func WithMaxBatchSize(size int) Option {
	return func(db *DB) error {
		db.maxSize = size
		return nil
	}
}

// DB represents the default AzmoDB backend. DB represents a collection of buckets
// persisted to a file on disk. All data access is performed through transactions
// which can be obtained through the DB.
type DB struct {
	root       *btree.DB
	maxEntries int
	maxSize    int
}

const (
	defaultMaxBatchSize    = 2 << 20
	defaultMaxBatchEntries = 256
)

// Open creates and opens a database at the given path. If the file does
// not exist then it will be created automatically.
//
// Timeout is the amount of time to wait to obtain a file lock. When set
// to zero it will wait indefinitely. This option is only available on
// Darwin and Linux.
func Open(path string, timeout time.Duration, opts ...Option) (*DB, error) {
	db := &DB{
		maxEntries: defaultMaxBatchEntries,
		maxSize:    defaultMaxBatchSize,
	}
	for _, opt := range opts {
		if err := opt(db); err != nil {
			return nil, err
		}
	}

	root, err := btree.Open(path, 0600, &btree.Options{
		Timeout: timeout,
	})
	if err != nil {
		return nil, err
	}

	if err = root.Update(func(tx *btree.Tx) (err error) {
		for _, name := range rootBuckets {
			_, err = tx.CreateBucketIfNotExists(name)
			if err != nil {
				break
			}
		}
		return err
	}); err != nil {
		return nil, err
	}

	db.root = root
	return db, nil
}

// Close releases all database resources. All batch transactions must be
// closed before closing the database.
func (db *DB) Close() error {
	if db == nil || db.root == nil {
		return errors.New("backend is shut down")
	}

	err := db.root.Close()
	db.root = nil
	return err
}

// WriteTo writes the entire database to a writer.
func (db *DB) WriteTo(w io.Writer) (n int64, err error) {
	err = db.root.View(func(tx *btree.Tx) error {
		n, err = tx.WriteTo(w)
		return err
	})
	return n, err
}

// Range performs fn on all values stored in the database at rev. If fn
// returns an errors the traversal is stopped and the error is returned.
func (db *DB) Range(rev Revision, fn func(key, value []byte) error) error {
	return db.root.View(func(tx *btree.Tx) (err error) {
		meta := tx.Bucket(metaBucket).Bucket(rev[:])
		if meta == nil {
			return errors.New("revision not found")
		}
		data := tx.Bucket(dataBucket)

		c := meta.Cursor()
		for k, sum := c.First(); k != nil; k, sum = c.Next() {
			v := data.Get(sum)
			if v == nil {
				panic("cannot find value for key: " + string(k))
			}
			safeKey := clone(nil, k)
			safeValue, err := snappy.Decode(nil, v)
			if err != nil {
				return err
			}
			if err = fn(safeKey, safeValue); err != nil {
				break
			}
		}
		return err
	})
}

// Last returns the last revision in the database and an error if any.
func (db *DB) Last() (rev Revision, err error) {
	err = db.root.View(func(tx *btree.Tx) error {
		c := tx.Bucket(metaBucket).Cursor()
		k, _ := c.Last()
		copy(rev[:], k)
		return nil
	})
	return rev, err
}

// Batch starts a new batch transaction. Starting multiple write batch
// transactions will cause the calls to block and be serialized until
// the current write batch transaction finishes.
func (db *DB) Batch(rev Revision) (Batch, error) {
	tx, err := db.root.Begin(true)
	if err != nil {
		return nil, err
	}

	meta, data := tx.Bucket(metaBucket), tx.Bucket(dataBucket)
	meta, err = meta.CreateBucket(rev[:])
	if err != nil {
		return nil, err
	}
	return &batch{
		entries:    make([]*entry, db.maxEntries),
		maxEntries: db.maxEntries,
		maxSize:    db.maxSize,
		meta:       meta,
		data:       data,
		tx:         tx,
	}, nil
}

func sha1sum(data []byte) [sha1.Size]byte {
	b := [sha1.Size]byte{}
	h := sha1.New()
	h.Write(data)
	h.Sum(b[:0])
	return b
}

type entry struct {
	key   []byte
	value []byte
}

type batch struct {
	entries    []*entry
	index      int
	size       int
	maxEntries int
	maxSize    int

	meta *btree.Bucket
	data *btree.Bucket
	tx   *btree.Tx
}

func (b *batch) next() *entry {
	e := b.entries[b.index]
	if e == nil {
		e = &entry{}
		b.entries[b.index] = e
	}
	return e
}

func (b *batch) Put(key, value []byte) error {
	e := b.next()
	e.key = clone(nil, key)
	e.value = snappy.Encode(nil, value)
	b.size += len(key) + len(value)
	b.index++

	return b.flush(false)
}

func (b *batch) put(key, value []byte) (err error) {
	sum := sha1sum(value)
	if v := b.data.Get(sum[:]); v == nil {
		err = b.data.Put(sum[:], value)
		if err != nil {
			return err
		}
	}
	return b.meta.Put(key, sum[:])
}

func (b *batch) flush(force bool) (err error) {
	if b.index >= b.maxEntries || b.size >= b.maxSize || force {
		for i := 0; i < b.index; i++ {
			e := b.entries[i]
			if err = b.put(e.key, e.value); err != nil {
				return err
			}
			e.key = nil
			e.value = nil
		}
		b.index = 0
		b.size = 0
	}
	return err
}

func (b *batch) Close() error {
	if err := b.flush(true); err != nil {
		//		b.tx.Rollback()
		return err
	}
	return b.tx.Commit()
}

func clone(dst, src []byte) []byte {
	n := len(src)
	if len(dst) < n {
		dst = make([]byte, n)
	}
	dst = dst[:n]
	copy(dst, src)
	return dst
}
<commit_msg>fix compress deadlock / default to compression<commit_after>// Package backend implements the persistent AzmoDB backend key/value
// database.
package backend

import (
	"crypto/sha1"
	"errors"
	"io"
	"time"

	btree "github.com/boltdb/bolt"
	"github.com/golang/snappy"
)

// Backend represents a persistent AzmoDB backend.
type Backend interface {
	// Range performs fn on all values stored in the database at rev.
	// If fn returns an errors the traversal is stopped and the error
	// is returned.
	Range(rev Revision, fn func(key []byte, value []byte) error) error

	// Batch starts a new batch transaction. Starting multiple write
	// batch transactions will cause the calls to block and be
	// serialized until the current write batch transaction finishes.
	Batch(rev Revision) (Batch, error)

	// Last returns the last revision in the database and an error if
	// any.
	Last() (Revision, error)
}

// Batch returns a batch transaction on the database.
type Batch interface {
	// Put sets the value for a key in the database. Put must create a
	// copy of the supplied key and value.
	Put(key []byte, value []byte) error

	// Close closes the batch transaction.
	Close() error
}

// Revision represents a serialized AzmoDB revision.
type Revision [8]byte

var (
	rootBuckets = [][]byte{dataBucket, metaBucket}
	dataBucket  = []byte("__data__")
	metaBucket  = []byte("__meta__")

	_ Backend = (*DB)(nil)
)

// Option represents a DB option function.
type Option func(*DB) error

// WithMaxBatchEntries configures the maximum batch entries.
func WithMaxBatchEntries(entries int) Option {
	return func(db *DB) error {
		db.maxEntries = entries
		return nil
	}
}

// WithMaxBatchSize configures the maximum batch size.
func WithMaxBatchSize(size int) Option {
	return func(db *DB) error {
		db.maxSize = size
		return nil
	}
}

// DB represents the default AzmoDB backend. DB represents a collection of buckets
// persisted to a file on disk. All data access is performed through transactions
// which can be obtained through the DB.
type DB struct {
	root       *btree.DB
	maxEntries int
	maxSize    int
}

const (
	defaultMaxBatchSize    = 2 << 20
	defaultMaxBatchEntries = 256
)

// Open creates and opens a database at the given path. If the file does
// not exist then it will be created automatically.
//
// Timeout is the amount of time to wait to obtain a file lock. When set
// to zero it will wait indefinitely. This option is only available on
// Darwin and Linux.
func Open(path string, timeout time.Duration, opts ...Option) (*DB, error) {
	db := &DB{
		maxEntries: defaultMaxBatchEntries,
		maxSize:    defaultMaxBatchSize,
	}
	for _, opt := range opts {
		if err := opt(db); err != nil {
			return nil, err
		}
	}

	root, err := btree.Open(path, 0600, &btree.Options{
		Timeout: timeout,
	})
	if err != nil {
		return nil, err
	}

	if err = root.Update(func(tx *btree.Tx) (err error) {
		for _, name := range rootBuckets {
			_, err = tx.CreateBucketIfNotExists(name)
			if err != nil {
				break
			}
		}
		return err
	}); err != nil {
		return nil, err
	}

	db.root = root
	return db, nil
}

// Close releases all database resources. All batch transactions must be
// closed before closing the database.
func (db *DB) Close() error {
	if db == nil || db.root == nil {
		return errors.New("backend is shut down")
	}

	err := db.root.Close()
	db.root = nil
	return err
}

// WriteTo writes the entire database to a writer.
func (db *DB) WriteTo(w io.Writer) (n int64, err error) {
	err = db.root.View(func(tx *btree.Tx) error {
		n, err = tx.WriteTo(w)
		return err
	})
	return n, err
}

// Range performs fn on all values stored in the database at rev. If fn
// returns an errors the traversal is stopped and the error is returned.
func (db *DB) Range(rev Revision, fn func(key, value []byte) error) error {
	return db.root.View(func(tx *btree.Tx) (err error) {
		meta := tx.Bucket(metaBucket).Bucket(rev[:])
		if meta == nil {
			return errors.New("revision not found")
		}
		data := tx.Bucket(dataBucket)

		c := meta.Cursor()
		for k, sum := c.First(); k != nil; k, sum = c.Next() {
			v := data.Get(sum)
			if v == nil {
				panic("cannot find value for key: " + string(k))
			}
			safeKey := clone(nil, k)
			safeValue, err := snappy.Decode(nil, v)
			if err != nil {
				return err
			}
			if err = fn(safeKey, safeValue); err != nil {
				break
			}
		}
		return err
	})
}

// Last returns the last revision in the database and an error if any.
func (db *DB) Last() (rev Revision, err error) {
	err = db.root.View(func(tx *btree.Tx) error {
		c := tx.Bucket(metaBucket).Cursor()
		k, _ := c.Last()
		copy(rev[:], k)
		return nil
	})
	return rev, err
}

// Batch starts a new batch transaction. Starting multiple write batch
// transactions will cause the calls to block and be serialized until
// the current write batch transaction finishes.
func (db *DB) Batch(rev Revision) (Batch, error) {
	tx, err := db.root.Begin(true)
	if err != nil {
		return nil, err
	}

	meta, data := tx.Bucket(metaBucket), tx.Bucket(dataBucket)
	meta, err = meta.CreateBucket(rev[:])
	if err != nil {
		return nil, err
	}
	return &batch{
		entries:    make([]*entry, db.maxEntries),
		maxEntries: db.maxEntries,
		maxSize:    db.maxSize,
		meta:       meta,
		data:       data,
		tx:         tx,
	}, nil
}

func sha1sum(data []byte) [sha1.Size]byte {
	b := [sha1.Size]byte{}
	h := sha1.New()
	h.Write(data)
	h.Sum(b[:0])
	return b
}

type entry struct {
	key   []byte
	value []byte
}

type batch struct {
	entries    []*entry
	index      int
	size       int
	maxEntries int
	maxSize    int

	meta *btree.Bucket
	data *btree.Bucket
	tx   *btree.Tx
}

func (b *batch) next() *entry {
	e := b.entries[b.index]
	if e == nil {
		e = &entry{}
		b.entries[b.index] = e
	}
	return e
}

func (b *batch) Put(key, value []byte) error {
	e := b.next()
	e.key = clone(nil, key)
	e.value = snappy.Encode(nil, value)
	b.size += len(key) + len(value)
	b.index++

	return b.flush(false)
}

func (b *batch) put(key, value []byte) (err error) {
	sum := sha1sum(value)
	if v := b.data.Get(sum[:]); v == nil {
		err = b.data.Put(sum[:], value)
		if err != nil {
			return err
		}
	}
	return b.meta.Put(key, sum[:])
}

func (b *batch) flush(force bool) (err error) {
	if b.index >= b.maxEntries || b.size >= b.maxSize || force {
		for i := 0; i < b.index; i++ {
			e := b.entries[i]
			if err = b.put(e.key, e.value); err != nil {
				return err
			}
			e.key = nil
			e.value = nil
		}
		b.index = 0
		b.size = 0
	}
	return err
}

func (b *batch) Close() error {
	if err := b.flush(true); err != nil {
		b.tx.Rollback()
		return err
	}
	return b.tx.Commit()
}

func clone(dst, src []byte) []byte {
	n := len(src)
	if len(dst) < n {
		dst = make([]byte, n)
	}
	dst = dst[:n]
	copy(dst, src)
	return dst
}
<|endoftext|>


<commit_before>// +build !windows

package main

import (
	"encoding/json"
	"fmt"
	"strings"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/mount"
	"github.com/docker/docker/api/types/swarm"
	"github.com/docker/docker/integration-cli/checker"
	"github.com/go-check/check"
)

func (s *DockerSwarmSuite) TestServiceCreateMountVolume(c *check.C) {
	d := s.AddDaemon(c, true, true)
	out, err := d.Cmd("service", "create", "--detach=true", "--mount", "type=volume,source=foo,target=/foo,volume-nocopy", "busybox", "top")
	c.Assert(err, checker.IsNil, check.Commentf(out))
	id := strings.TrimSpace(out)

	var tasks []swarm.Task
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		tasks = d.GetServiceTasks(c, id)
		return len(tasks) > 0, nil
	}, checker.Equals, true)

	task := tasks[0]
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		if task.NodeID == "" || task.Status.ContainerStatus.ContainerID == "" {
			task = d.GetTask(c, task.ID)
		}
		return task.NodeID != "" && task.Status.ContainerStatus.ContainerID != "", nil
	}, checker.Equals, true)

	// check container mount config
	out, err = s.nodeCmd(c, task.NodeID, "inspect", "--format", "{{json .HostConfig.Mounts}}", task.Status.ContainerStatus.ContainerID)
	c.Assert(err, checker.IsNil, check.Commentf(out))

	var mountConfig []mount.Mount
	c.Assert(json.Unmarshal([]byte(out), &mountConfig), checker.IsNil)
	c.Assert(mountConfig, checker.HasLen, 1)

	c.Assert(mountConfig[0].Source, checker.Equals, "foo")
	c.Assert(mountConfig[0].Target, checker.Equals, "/foo")
	c.Assert(mountConfig[0].Type, checker.Equals, mount.TypeVolume)
	c.Assert(mountConfig[0].VolumeOptions, checker.NotNil)
	c.Assert(mountConfig[0].VolumeOptions.NoCopy, checker.True)

	// check container mounts actual
	out, err = s.nodeCmd(c, task.NodeID, "inspect", "--format", "{{json .Mounts}}", task.Status.ContainerStatus.ContainerID)
	c.Assert(err, checker.IsNil, check.Commentf(out))

	var mounts []types.MountPoint
	c.Assert(json.Unmarshal([]byte(out), &mounts), checker.IsNil)
	c.Assert(mounts, checker.HasLen, 1)

	c.Assert(mounts[0].Type, checker.Equals, mount.TypeVolume)
	c.Assert(mounts[0].Name, checker.Equals, "foo")
	c.Assert(mounts[0].Destination, checker.Equals, "/foo")
	c.Assert(mounts[0].RW, checker.Equals, true)
}

func (s *DockerSwarmSuite) TestServiceCreateWithSecretSimple(c *check.C) {
	d := s.AddDaemon(c, true, true)

	serviceName := "test-service-secret"
	testName := "test_secret"
	id := d.CreateSecret(c, swarm.SecretSpec{
		Annotations: swarm.Annotations{
			Name: testName,
		},
		Data: []byte("TESTINGDATA"),
	})
	c.Assert(id, checker.Not(checker.Equals), "", check.Commentf("secrets: %s", id))

	out, err := d.Cmd("service", "create", "--name", serviceName, "--secret", testName, "busybox", "top")
	c.Assert(err, checker.IsNil, check.Commentf(out))

	out, err = d.Cmd("service", "inspect", "--format", "{{ json .Spec.TaskTemplate.ContainerSpec.Secrets }}", serviceName)
	c.Assert(err, checker.IsNil)

	var refs []swarm.SecretReference
	c.Assert(json.Unmarshal([]byte(out), &refs), checker.IsNil)
	c.Assert(refs, checker.HasLen, 1)

	c.Assert(refs[0].SecretName, checker.Equals, testName)
	c.Assert(refs[0].File, checker.Not(checker.IsNil))
	c.Assert(refs[0].File.Name, checker.Equals, testName)
	c.Assert(refs[0].File.UID, checker.Equals, "0")
	c.Assert(refs[0].File.GID, checker.Equals, "0")

	out, err = d.Cmd("service", "rm", serviceName)
	c.Assert(err, checker.IsNil, check.Commentf(out))
	d.DeleteSecret(c, testName)
}

func (s *DockerSwarmSuite) TestServiceCreateWithSecretSourceTargetPaths(c *check.C) {
	d := s.AddDaemon(c, true, true)

	testPaths := map[string]string{
		"app":         "/etc/secret",
		"test_secret": "test_secret",
	}
	for testName, testTarget := range testPaths {
		serviceName := "svc-" + testName
		id := d.CreateSecret(c, swarm.SecretSpec{
			Annotations: swarm.Annotations{
				Name: testName,
			},
			Data: []byte("TESTINGDATA"),
		})
		c.Assert(id, checker.Not(checker.Equals), "", check.Commentf("secrets: %s", id))

		out, err := d.Cmd("service", "create", "--name", serviceName, "--secret", fmt.Sprintf("source=%s,target=%s", testName, testTarget), "busybox", "top")
		c.Assert(err, checker.IsNil, check.Commentf(out))

		out, err = d.Cmd("service", "inspect", "--format", "{{ json .Spec.TaskTemplate.ContainerSpec.Secrets }}", serviceName)
		c.Assert(err, checker.IsNil)

		var refs []swarm.SecretReference
		c.Assert(json.Unmarshal([]byte(out), &refs), checker.IsNil)
		c.Assert(refs, checker.HasLen, 1)

		c.Assert(refs[0].SecretName, checker.Equals, testName)
		c.Assert(refs[0].File, checker.Not(checker.IsNil))
		c.Assert(refs[0].File.Name, checker.Equals, testTarget)

		out, err = d.Cmd("service", "rm", serviceName)
		c.Assert(err, checker.IsNil, check.Commentf(out))

		d.DeleteSecret(c, testName)
	}
}

func (s *DockerSwarmSuite) TestServiceCreateMountTmpfs(c *check.C) {
	d := s.AddDaemon(c, true, true)
	out, err := d.Cmd("service", "create", "--detach=true", "--mount", "type=tmpfs,target=/foo,tmpfs-size=1MB", "busybox", "sh", "-c", "mount | grep foo; tail -f /dev/null")
	c.Assert(err, checker.IsNil, check.Commentf(out))
	id := strings.TrimSpace(out)

	var tasks []swarm.Task
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		tasks = d.GetServiceTasks(c, id)
		return len(tasks) > 0, nil
	}, checker.Equals, true)

	task := tasks[0]
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		if task.NodeID == "" || task.Status.ContainerStatus.ContainerID == "" {
			task = d.GetTask(c, task.ID)
		}
		return task.NodeID != "" && task.Status.ContainerStatus.ContainerID != "", nil
	}, checker.Equals, true)

	// check container mount config
	out, err = s.nodeCmd(c, task.NodeID, "inspect", "--format", "{{json .HostConfig.Mounts}}", task.Status.ContainerStatus.ContainerID)
	c.Assert(err, checker.IsNil, check.Commentf(out))

	var mountConfig []mount.Mount
	c.Assert(json.Unmarshal([]byte(out), &mountConfig), checker.IsNil)
	c.Assert(mountConfig, checker.HasLen, 1)

	c.Assert(mountConfig[0].Source, checker.Equals, "")
	c.Assert(mountConfig[0].Target, checker.Equals, "/foo")
	c.Assert(mountConfig[0].Type, checker.Equals, mount.TypeTmpfs)
	c.Assert(mountConfig[0].TmpfsOptions, checker.NotNil)
	c.Assert(mountConfig[0].TmpfsOptions.SizeBytes, checker.Equals, int64(1048576))

	// check container mounts actual
	out, err = s.nodeCmd(c, task.NodeID, "inspect", "--format", "{{json .Mounts}}", task.Status.ContainerStatus.ContainerID)
	c.Assert(err, checker.IsNil, check.Commentf(out))

	var mounts []types.MountPoint
	c.Assert(json.Unmarshal([]byte(out), &mounts), checker.IsNil)
	c.Assert(mounts, checker.HasLen, 1)

	c.Assert(mounts[0].Type, checker.Equals, mount.TypeTmpfs)
	c.Assert(mounts[0].Name, checker.Equals, "")
	c.Assert(mounts[0].Destination, checker.Equals, "/foo")
	c.Assert(mounts[0].RW, checker.Equals, true)

	out, err = s.nodeCmd(c, task.NodeID, "logs", task.Status.ContainerStatus.ContainerID)
	c.Assert(err, checker.IsNil, check.Commentf(out))
	c.Assert(strings.TrimSpace(out), checker.HasPrefix, "tmpfs on /foo type tmpfs")
	c.Assert(strings.TrimSpace(out), checker.Contains, "size=1024k")
}
<commit_msg>Extend test coverage of secrets<commit_after>// +build !windows

package main

import (
	"encoding/json"
	"fmt"
	"path/filepath"
	"strings"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/mount"
	"github.com/docker/docker/api/types/swarm"
	"github.com/docker/docker/integration-cli/checker"
	"github.com/go-check/check"
)

func (s *DockerSwarmSuite) TestServiceCreateMountVolume(c *check.C) {
	d := s.AddDaemon(c, true, true)
	out, err := d.Cmd("service", "create", "--detach=true", "--mount", "type=volume,source=foo,target=/foo,volume-nocopy", "busybox", "top")
	c.Assert(err, checker.IsNil, check.Commentf(out))
	id := strings.TrimSpace(out)

	var tasks []swarm.Task
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		tasks = d.GetServiceTasks(c, id)
		return len(tasks) > 0, nil
	}, checker.Equals, true)

	task := tasks[0]
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		if task.NodeID == "" || task.Status.ContainerStatus.ContainerID == "" {
			task = d.GetTask(c, task.ID)
		}
		return task.NodeID != "" && task.Status.ContainerStatus.ContainerID != "", nil
	}, checker.Equals, true)

	// check container mount config
	out, err = s.nodeCmd(c, task.NodeID, "inspect", "--format", "{{json .HostConfig.Mounts}}", task.Status.ContainerStatus.ContainerID)
	c.Assert(err, checker.IsNil, check.Commentf(out))

	var mountConfig []mount.Mount
	c.Assert(json.Unmarshal([]byte(out), &mountConfig), checker.IsNil)
	c.Assert(mountConfig, checker.HasLen, 1)

	c.Assert(mountConfig[0].Source, checker.Equals, "foo")
	c.Assert(mountConfig[0].Target, checker.Equals, "/foo")
	c.Assert(mountConfig[0].Type, checker.Equals, mount.TypeVolume)
	c.Assert(mountConfig[0].VolumeOptions, checker.NotNil)
	c.Assert(mountConfig[0].VolumeOptions.NoCopy, checker.True)

	// check container mounts actual
	out, err = s.nodeCmd(c, task.NodeID, "inspect", "--format", "{{json .Mounts}}", task.Status.ContainerStatus.ContainerID)
	c.Assert(err, checker.IsNil, check.Commentf(out))

	var mounts []types.MountPoint
	c.Assert(json.Unmarshal([]byte(out), &mounts), checker.IsNil)
	c.Assert(mounts, checker.HasLen, 1)

	c.Assert(mounts[0].Type, checker.Equals, mount.TypeVolume)
	c.Assert(mounts[0].Name, checker.Equals, "foo")
	c.Assert(mounts[0].Destination, checker.Equals, "/foo")
	c.Assert(mounts[0].RW, checker.Equals, true)
}

func (s *DockerSwarmSuite) TestServiceCreateWithSecretSimple(c *check.C) {
	d := s.AddDaemon(c, true, true)

	serviceName := "test-service-secret"
	testName := "test_secret"
	id := d.CreateSecret(c, swarm.SecretSpec{
		Annotations: swarm.Annotations{
			Name: testName,
		},
		Data: []byte("TESTINGDATA"),
	})
	c.Assert(id, checker.Not(checker.Equals), "", check.Commentf("secrets: %s", id))

	out, err := d.Cmd("service", "create", "--name", serviceName, "--secret", testName, "busybox", "top")
	c.Assert(err, checker.IsNil, check.Commentf(out))

	out, err = d.Cmd("service", "inspect", "--format", "{{ json .Spec.TaskTemplate.ContainerSpec.Secrets }}", serviceName)
	c.Assert(err, checker.IsNil)

	var refs []swarm.SecretReference
	c.Assert(json.Unmarshal([]byte(out), &refs), checker.IsNil)
	c.Assert(refs, checker.HasLen, 1)

	c.Assert(refs[0].SecretName, checker.Equals, testName)
	c.Assert(refs[0].File, checker.Not(checker.IsNil))
	c.Assert(refs[0].File.Name, checker.Equals, testName)
	c.Assert(refs[0].File.UID, checker.Equals, "0")
	c.Assert(refs[0].File.GID, checker.Equals, "0")

	out, err = d.Cmd("service", "rm", serviceName)
	c.Assert(err, checker.IsNil, check.Commentf(out))
	d.DeleteSecret(c, testName)
}

func (s *DockerSwarmSuite) TestServiceCreateWithSecretSourceTargetPaths(c *check.C) {
	d := s.AddDaemon(c, true, true)

	testPaths := map[string]string{
		"app":                  "/etc/secret",
		"test_secret":          "test_secret",
		"relative_secret":      "relative/secret",
		"escapes_in_container": "../secret",
	}

	var secretFlags []string

	for testName, testTarget := range testPaths {
		id := d.CreateSecret(c, swarm.SecretSpec{
			Annotations: swarm.Annotations{
				Name: testName,
			},
			Data: []byte("TESTINGDATA " + testName + " " + testTarget),
		})
		c.Assert(id, checker.Not(checker.Equals), "", check.Commentf("secrets: %s", id))

		secretFlags = append(secretFlags, "--secret", fmt.Sprintf("source=%s,target=%s", testName, testTarget))
	}

	serviceName := "svc"
	serviceCmd := []string{"service", "create", "--name", serviceName}
	serviceCmd = append(serviceCmd, secretFlags...)
	serviceCmd = append(serviceCmd, "busybox", "top")
	out, err := d.Cmd(serviceCmd...)
	c.Assert(err, checker.IsNil, check.Commentf(out))

	out, err = d.Cmd("service", "inspect", "--format", "{{ json .Spec.TaskTemplate.ContainerSpec.Secrets }}", serviceName)
	c.Assert(err, checker.IsNil)

	var refs []swarm.SecretReference
	c.Assert(json.Unmarshal([]byte(out), &refs), checker.IsNil)
	c.Assert(refs, checker.HasLen, len(testPaths))

	var tasks []swarm.Task
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		tasks = d.GetServiceTasks(c, serviceName)
		return len(tasks) > 0, nil
	}, checker.Equals, true)

	task := tasks[0]
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		if task.NodeID == "" || task.Status.ContainerStatus.ContainerID == "" {
			task = d.GetTask(c, task.ID)
		}
		return task.NodeID != "" && task.Status.ContainerStatus.ContainerID != "", nil
	}, checker.Equals, true)

	for testName, testTarget := range testPaths {
		path := testTarget
		if !filepath.IsAbs(path) {
			path = filepath.Join("/run/secrets", path)
		}
		out, err := d.Cmd("exec", task.Status.ContainerStatus.ContainerID, "cat", path)
		c.Assert(err, checker.IsNil)
		c.Assert(out, checker.Equals, "TESTINGDATA "+testName+" "+testTarget)
	}

	out, err = d.Cmd("service", "rm", serviceName)
	c.Assert(err, checker.IsNil, check.Commentf(out))
}

func (s *DockerSwarmSuite) TestServiceCreateWithSecretReferencedTwice(c *check.C) {
	d := s.AddDaemon(c, true, true)

	id := d.CreateSecret(c, swarm.SecretSpec{
		Annotations: swarm.Annotations{
			Name: "mysecret",
		},
		Data: []byte("TESTINGDATA"),
	})
	c.Assert(id, checker.Not(checker.Equals), "", check.Commentf("secrets: %s", id))

	serviceName := "svc"
	out, err := d.Cmd("service", "create", "--name", serviceName, "--secret", "source=mysecret,target=target1", "--secret", "source=mysecret,target=target2", "busybox", "top")
	c.Assert(err, checker.IsNil, check.Commentf(out))

	out, err = d.Cmd("service", "inspect", "--format", "{{ json .Spec.TaskTemplate.ContainerSpec.Secrets }}", serviceName)
	c.Assert(err, checker.IsNil)

	var refs []swarm.SecretReference
	c.Assert(json.Unmarshal([]byte(out), &refs), checker.IsNil)
	c.Assert(refs, checker.HasLen, 2)

	var tasks []swarm.Task
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		tasks = d.GetServiceTasks(c, serviceName)
		return len(tasks) > 0, nil
	}, checker.Equals, true)

	task := tasks[0]
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		if task.NodeID == "" || task.Status.ContainerStatus.ContainerID == "" {
			task = d.GetTask(c, task.ID)
		}
		return task.NodeID != "" && task.Status.ContainerStatus.ContainerID != "", nil
	}, checker.Equals, true)

	for _, target := range []string{"target1", "target2"} {
		c.Assert(err, checker.IsNil, check.Commentf(out))
		path := filepath.Join("/run/secrets", target)
		out, err := d.Cmd("exec", task.Status.ContainerStatus.ContainerID, "cat", path)
		c.Assert(err, checker.IsNil)
		c.Assert(out, checker.Equals, "TESTINGDATA")
	}

	out, err = d.Cmd("service", "rm", serviceName)
	c.Assert(err, checker.IsNil, check.Commentf(out))
}

func (s *DockerSwarmSuite) TestServiceCreateMountTmpfs(c *check.C) {
	d := s.AddDaemon(c, true, true)
	out, err := d.Cmd("service", "create", "--detach=true", "--mount", "type=tmpfs,target=/foo,tmpfs-size=1MB", "busybox", "sh", "-c", "mount | grep foo; tail -f /dev/null")
	c.Assert(err, checker.IsNil, check.Commentf(out))
	id := strings.TrimSpace(out)

	var tasks []swarm.Task
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		tasks = d.GetServiceTasks(c, id)
		return len(tasks) > 0, nil
	}, checker.Equals, true)

	task := tasks[0]
	waitAndAssert(c, defaultReconciliationTimeout, func(c *check.C) (interface{}, check.CommentInterface) {
		if task.NodeID == "" || task.Status.ContainerStatus.ContainerID == "" {
			task = d.GetTask(c, task.ID)
		}
		return task.NodeID != "" && task.Status.ContainerStatus.ContainerID != "", nil
	}, checker.Equals, true)

	// check container mount config
	out, err = s.nodeCmd(c, task.NodeID, "inspect", "--format", "{{json .HostConfig.Mounts}}", task.Status.ContainerStatus.ContainerID)
	c.Assert(err, checker.IsNil, check.Commentf(out))

	var mountConfig []mount.Mount
	c.Assert(json.Unmarshal([]byte(out), &mountConfig), checker.IsNil)
	c.Assert(mountConfig, checker.HasLen, 1)

	c.Assert(mountConfig[0].Source, checker.Equals, "")
	c.Assert(mountConfig[0].Target, checker.Equals, "/foo")
	c.Assert(mountConfig[0].Type, checker.Equals, mount.TypeTmpfs)
	c.Assert(mountConfig[0].TmpfsOptions, checker.NotNil)
	c.Assert(mountConfig[0].TmpfsOptions.SizeBytes, checker.Equals, int64(1048576))

	// check container mounts actual
	out, err = s.nodeCmd(c, task.NodeID, "inspect", "--format", "{{json .Mounts}}", task.Status.ContainerStatus.ContainerID)
	c.Assert(err, checker.IsNil, check.Commentf(out))

	var mounts []types.MountPoint
	c.Assert(json.Unmarshal([]byte(out), &mounts), checker.IsNil)
	c.Assert(mounts, checker.HasLen, 1)

	c.Assert(mounts[0].Type, checker.Equals, mount.TypeTmpfs)
	c.Assert(mounts[0].Name, checker.Equals, "")
	c.Assert(mounts[0].Destination, checker.Equals, "/foo")
	c.Assert(mounts[0].RW, checker.Equals, true)

	out, err = s.nodeCmd(c, task.NodeID, "logs", task.Status.ContainerStatus.ContainerID)
	c.Assert(err, checker.IsNil, check.Commentf(out))
	c.Assert(strings.TrimSpace(out), checker.HasPrefix, "tmpfs on /foo type tmpfs")
	c.Assert(strings.TrimSpace(out), checker.Contains, "size=1024k")
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/grsakea/kappastat/common"
	"github.com/mrshankly/go-twitch/twitch"
	"log"
	"time"
)

func loopViewers(client *twitch.Client, c chan Message, infos chan kappastat.ViewerCount) {
	followed := []string{}
	ticker := time.NewTicker(time.Minute).C

	for {
		select {
		case msg := <-c:
			followed = followedHandler(followed, msg)
		case <-ticker:
			for _, v := range followed {
				infos <- fetchViewers(client, v)
			}
		}
	}
}

func fetchViewers(client *twitch.Client, chan_string string) kappastat.ViewerCount {

	channel, err := client.Streams.Channel(chan_string)
	if err != nil {
		channel, err = client.Streams.Channel(chan_string)
		if err != nil {
			log.Fatal(err)
		}
	}

	return kappastat.ViewerCount{chan_string, time.Now(), channel.Stream.Viewers}
}
<commit_msg>Stop an api error from crashing the program<commit_after>package main

import (
	"github.com/grsakea/kappastat/common"
	"github.com/mrshankly/go-twitch/twitch"
	"log"
	"time"
)

func loopViewers(client *twitch.Client, c chan Message, infos chan kappastat.ViewerCount) {
	followed := []string{}
	ticker := time.NewTicker(time.Minute).C

	for {
		select {
		case msg := <-c:
			followed = followedHandler(followed, msg)
		case <-ticker:
			for _, v := range followed {
				infos <- fetchViewers(client, v)
			}
		}
	}
}

func fetchViewers(client *twitch.Client, chan_string string) kappastat.ViewerCount {

	channel, err := client.Streams.Channel(chan_string)
	if err != nil {
		channel, err = client.Streams.Channel(chan_string)
		if err != nil {
			log.Print(err)
		}
	}

	return kappastat.ViewerCount{chan_string, time.Now(), channel.Stream.Viewers}
}
<|endoftext|>


<commit_before>package kcp

import (
	"sync"

	"github.com/v2ray/v2ray-core/common/alloc"
)

type SendingWindow struct {
	start uint32
	cap   uint32
	len   uint32
	last  uint32

	data []*DataSegment
	prev []uint32
	next []uint32

	inFlightSize uint32
	writer       SegmentWriter
	onPacketLoss func(bool)
}

func NewSendingWindow(size uint32, inFlightSize uint32, writer SegmentWriter, onPacketLoss func(bool)) *SendingWindow {
	window := &SendingWindow{
		start:        0,
		cap:          size,
		len:          0,
		last:         0,
		data:         make([]*DataSegment, size),
		prev:         make([]uint32, size),
		next:         make([]uint32, size),
		writer:       writer,
		onPacketLoss: onPacketLoss,
		inFlightSize: inFlightSize,
	}
	return window
}

func (this *SendingWindow) Len() int {
	return int(this.len)
}

func (this *SendingWindow) Push(seg *DataSegment) {
	pos := (this.start + this.len) % this.cap
	this.data[pos] = seg
	if this.len > 0 {
		this.next[this.last] = pos
		this.prev[pos] = this.last
	}
	this.last = pos
	this.len++
}

func (this *SendingWindow) First() *DataSegment {
	return this.data[this.start]
}

func (this *SendingWindow) Clear(una uint32) {
	for this.Len() > 0 && this.data[this.start].Number < una {
		this.Remove(0)
	}
}

func (this *SendingWindow) Remove(idx uint32) {
	if this.len == 0 {
		return
	}

	pos := (this.start + idx) % this.cap
	seg := this.data[pos]
	if seg == nil {
		return
	}
	seg.Release()
	this.data[pos] = nil
	if pos == this.start && pos == this.last {
		this.len = 0
		this.start = 0
		this.last = 0
	} else if pos == this.start {
		delta := this.next[pos] - this.start
		if this.next[pos] < this.start {
			delta = this.next[pos] + this.cap - this.start
		}
		this.start = this.next[pos]
		this.len -= delta
	} else if pos == this.last {
		this.last = this.prev[pos]
	} else {
		this.next[this.prev[pos]] = this.next[pos]
		this.prev[this.next[pos]] = this.prev[pos]
	}
}

func (this *SendingWindow) HandleFastAck(number uint32) {
	if this.len == 0 {
		return
	}

	for i := this.start; ; i = this.next[i] {
		seg := this.data[i]
		if _itimediff(number, seg.Number) < 0 {
			break
		}
		if number != seg.Number {
			seg.ackSkipped++
		}
		if i == this.last {
			break
		}
	}
}

func (this *SendingWindow) Flush(current uint32, resend uint32, rto uint32) {
	if this.Len() == 0 {
		return
	}

	lost := false
	var inFlightSize uint32

	for i := this.start; ; i = this.next[i] {
		segment := this.data[i]
		needsend := false
		if segment.transmit == 0 {
			needsend = true
			segment.transmit++
			segment.timeout = current + rto
		} else if _itimediff(current, segment.timeout) >= 0 {
			needsend = true
			segment.transmit++
			segment.timeout = current + rto
			lost = true
		} else if segment.ackSkipped >= resend {
			needsend = true
			segment.transmit++
			segment.ackSkipped = 0
			segment.timeout = current + rto
			lost = true
		}

		if needsend {
			this.writer.Write(segment)
			inFlightSize++
			if inFlightSize >= this.inFlightSize {
				break
			}
		}
		if i == this.last {
			break
		}
	}

	this.onPacketLoss(lost)
}

type SendingQueue struct {
	start uint32
	cap   uint32
	len   uint32
	list  []*DataSegment
}

func NewSendingQueue(size uint32) *SendingQueue {
	return &SendingQueue{
		start: 0,
		cap:   size,
		list:  make([]*DataSegment, size),
		len:   0,
	}
}

func (this *SendingQueue) IsFull() bool {
	return this.len == this.cap
}

func (this *SendingQueue) IsEmpty() bool {
	return this.len == 0
}

func (this *SendingQueue) Pop() *DataSegment {
	if this.IsEmpty() {
		return nil
	}
	seg := this.list[this.start]
	this.list[this.start] = nil
	this.len--
	this.start++
	if this.start == this.cap {
		this.start = 0
	}
	return seg
}

func (this *SendingQueue) Push(seg *DataSegment) {
	if this.IsFull() {
		return
	}
	this.list[(this.start+this.len)%this.cap] = seg
	this.len++
}

func (this *SendingQueue) Clear() {
	for i := uint32(0); i < this.len; i++ {
		this.list[(i+this.start)%this.cap].Release()
		this.list[(i+this.start)%this.cap] = nil
	}
	this.start = 0
	this.len = 0
}

func (this *SendingQueue) Len() uint32 {
	return this.len
}

type SendingWorker struct {
	sync.Mutex
	kcp                 *KCP
	window              *SendingWindow
	queue               *SendingQueue
	windowSize          uint32
	firstUnacknowledged uint32
	nextNumber          uint32
	remoteNextNumber    uint32
	controlWindow       uint32
	fastResend          uint32
	updated             bool
}

func NewSendingWorker(kcp *KCP) *SendingWorker {
	worker := &SendingWorker{
		kcp:              kcp,
		queue:            NewSendingQueue(effectiveConfig.GetSendingQueueSize()),
		fastResend:       2,
		remoteNextNumber: 32,
		windowSize:       effectiveConfig.GetSendingWindowSize(),
		controlWindow:    effectiveConfig.GetSendingWindowSize(),
	}
	worker.window = NewSendingWindow(effectiveConfig.GetSendingWindowSize(), effectiveConfig.GetSendingInFlightSize(), worker, worker.OnPacketLoss)
	return worker
}

func (this *SendingWorker) ProcessReceivingNext(nextNumber uint32) {
	this.Lock()
	defer this.Unlock()

	this.window.Clear(nextNumber)
	this.FindFirstUnacknowledged()
}

// @Private
func (this *SendingWorker) FindFirstUnacknowledged() {
	prevUna := this.firstUnacknowledged
	if this.window.Len() > 0 {
		this.firstUnacknowledged = this.window.First().Number
	} else {
		this.firstUnacknowledged = this.nextNumber
	}
	if this.firstUnacknowledged != prevUna {
		this.updated = true
	}
}

func (this *SendingWorker) ProcessAck(number uint32) {
	if number-this.firstUnacknowledged > this.windowSize {
		return
	}

	this.Lock()
	defer this.Unlock()
	this.window.Remove(number - this.firstUnacknowledged)
	this.FindFirstUnacknowledged()
}

func (this *SendingWorker) ProcessSegment(seg *AckSegment) {
	if this.remoteNextNumber < seg.ReceivingWindow {
		this.remoteNextNumber = seg.ReceivingWindow
	}
	this.ProcessReceivingNext(seg.ReceivingNext)
	var maxack uint32
	for i := 0; i < int(seg.Count); i++ {
		timestamp := seg.TimestampList[i]
		number := seg.NumberList[i]
		if this.kcp.current-timestamp > 10000 {
			this.kcp.update_ack(int32(this.kcp.current - timestamp))
		}
		this.ProcessAck(number)
		if maxack < number {
			maxack = number
		}
	}
	this.Lock()
	this.window.HandleFastAck(maxack)
	this.Unlock()
}

func (this *SendingWorker) Push(b []byte) int {
	nBytes := 0
	for len(b) > 0 && !this.queue.IsFull() {
		var size int
		if len(b) > int(this.kcp.mss) {
			size = int(this.kcp.mss)
		} else {
			size = len(b)
		}
		seg := &DataSegment{
			Data: alloc.NewSmallBuffer().Clear().Append(b[:size]),
		}
		this.Lock()
		this.queue.Push(seg)
		this.Unlock()
		b = b[size:]
		nBytes += size
	}
	return nBytes
}

func (this *SendingWorker) Write(seg Segment) {
	dataSeg := seg.(*DataSegment)

	dataSeg.Conv = this.kcp.conv
	dataSeg.Timestamp = this.kcp.current
	dataSeg.SendingNext = this.firstUnacknowledged
	dataSeg.Opt = 0
	if this.kcp.state == StateReadyToClose {
		dataSeg.Opt = SegmentOptionClose
	}

	this.kcp.output.Write(dataSeg)
	this.updated = false
}

func (this *SendingWorker) PingNecessary() bool {
	return this.updated
}

func (this *SendingWorker) OnPacketLoss(lost bool) {
	if !effectiveConfig.Congestion {
		return
	}

	if lost {
		this.controlWindow = 3 * this.controlWindow / 4
	} else {
		this.controlWindow += this.controlWindow / 4
	}
	if this.controlWindow < 4 {
		this.controlWindow = 4
	}
	if this.controlWindow > 2*this.windowSize {
		this.controlWindow = 2 * this.windowSize
	}
}

func (this *SendingWorker) Flush() {
	this.Lock()
	defer this.Unlock()

	cwnd := this.firstUnacknowledged + this.windowSize
	if cwnd > this.remoteNextNumber {
		cwnd = this.remoteNextNumber
	}
	if effectiveConfig.Congestion && cwnd > this.firstUnacknowledged+this.controlWindow {
		cwnd = this.firstUnacknowledged + this.controlWindow
	}

	for !this.queue.IsEmpty() && _itimediff(this.nextNumber, cwnd) < 0 {
		seg := this.queue.Pop()
		seg.Number = this.nextNumber
		seg.timeout = this.kcp.current
		seg.ackSkipped = 0
		seg.transmit = 0
		this.window.Push(seg)
		this.nextNumber++
	}

	this.window.Flush(this.kcp.current, this.kcp.fastresend, this.kcp.rx_rto)
}

func (this *SendingWorker) CloseWrite() {
	this.Lock()
	defer this.Unlock()

	this.window.Clear(0xFFFFFFFF)
	this.queue.Clear()
}
<commit_msg>refine congestion control<commit_after>package kcp

import (
	"sync"

	"github.com/v2ray/v2ray-core/common/alloc"
)

type SendingWindow struct {
	start uint32
	cap   uint32
	len   uint32
	last  uint32

	data []*DataSegment
	prev []uint32
	next []uint32

	inFlightSize      uint32
	totalInFlightSize uint32
	writer            SegmentWriter
	onPacketLoss      func(uint32)
}

func NewSendingWindow(size uint32, inFlightSize uint32, writer SegmentWriter, onPacketLoss func(uint32)) *SendingWindow {
	window := &SendingWindow{
		start:        0,
		cap:          size,
		len:          0,
		last:         0,
		data:         make([]*DataSegment, size),
		prev:         make([]uint32, size),
		next:         make([]uint32, size),
		writer:       writer,
		onPacketLoss: onPacketLoss,
		inFlightSize: inFlightSize,
	}
	return window
}

func (this *SendingWindow) Len() int {
	return int(this.len)
}

func (this *SendingWindow) Push(seg *DataSegment) {
	pos := (this.start + this.len) % this.cap
	this.data[pos] = seg
	if this.len > 0 {
		this.next[this.last] = pos
		this.prev[pos] = this.last
	}
	this.last = pos
	this.len++
}

func (this *SendingWindow) First() *DataSegment {
	return this.data[this.start]
}

func (this *SendingWindow) Clear(una uint32) {
	for this.Len() > 0 && this.data[this.start].Number < una {
		this.Remove(0)
	}
}

func (this *SendingWindow) Remove(idx uint32) {
	if this.len == 0 {
		return
	}

	pos := (this.start + idx) % this.cap
	seg := this.data[pos]
	if seg == nil {
		return
	}
	this.totalInFlightSize--
	seg.Release()
	this.data[pos] = nil
	if pos == this.start && pos == this.last {
		this.len = 0
		this.start = 0
		this.last = 0
	} else if pos == this.start {
		delta := this.next[pos] - this.start
		if this.next[pos] < this.start {
			delta = this.next[pos] + this.cap - this.start
		}
		this.start = this.next[pos]
		this.len -= delta
	} else if pos == this.last {
		this.last = this.prev[pos]
	} else {
		this.next[this.prev[pos]] = this.next[pos]
		this.prev[this.next[pos]] = this.prev[pos]
	}
}

func (this *SendingWindow) HandleFastAck(number uint32) {
	if this.len == 0 {
		return
	}

	for i := this.start; ; i = this.next[i] {
		seg := this.data[i]
		if _itimediff(number, seg.Number) < 0 {
			break
		}
		if number != seg.Number {
			seg.ackSkipped++
		}
		if i == this.last {
			break
		}
	}
}

func (this *SendingWindow) Flush(current uint32, resend uint32, rto uint32) {
	if this.Len() == 0 {
		return
	}

	var lost uint32
	var inFlightSize uint32

	for i := this.start; ; i = this.next[i] {
		segment := this.data[i]
		needsend := false
		if segment.transmit == 0 {
			needsend = true
			segment.transmit++
			segment.timeout = current + rto
			this.totalInFlightSize++
		} else if _itimediff(current, segment.timeout) >= 0 {
			needsend = true
			segment.transmit++
			segment.timeout = current + rto
			lost++
		} else if segment.ackSkipped >= resend {
			needsend = true
			segment.transmit++
			segment.ackSkipped = 0
			segment.timeout = current + rto
		}

		if needsend {
			this.writer.Write(segment)
			inFlightSize++
			if inFlightSize >= this.inFlightSize {
				break
			}
		}
		if i == this.last {
			break
		}
	}

	if inFlightSize > 0 && this.totalInFlightSize != 0 {
		rate := lost * 100 / this.totalInFlightSize
		this.onPacketLoss(rate)
	}
}

type SendingQueue struct {
	start uint32
	cap   uint32
	len   uint32
	list  []*DataSegment
}

func NewSendingQueue(size uint32) *SendingQueue {
	return &SendingQueue{
		start: 0,
		cap:   size,
		list:  make([]*DataSegment, size),
		len:   0,
	}
}

func (this *SendingQueue) IsFull() bool {
	return this.len == this.cap
}

func (this *SendingQueue) IsEmpty() bool {
	return this.len == 0
}

func (this *SendingQueue) Pop() *DataSegment {
	if this.IsEmpty() {
		return nil
	}
	seg := this.list[this.start]
	this.list[this.start] = nil
	this.len--
	this.start++
	if this.start == this.cap {
		this.start = 0
	}
	return seg
}

func (this *SendingQueue) Push(seg *DataSegment) {
	if this.IsFull() {
		return
	}
	this.list[(this.start+this.len)%this.cap] = seg
	this.len++
}

func (this *SendingQueue) Clear() {
	for i := uint32(0); i < this.len; i++ {
		this.list[(i+this.start)%this.cap].Release()
		this.list[(i+this.start)%this.cap] = nil
	}
	this.start = 0
	this.len = 0
}

func (this *SendingQueue) Len() uint32 {
	return this.len
}

type SendingWorker struct {
	sync.Mutex
	kcp                 *KCP
	window              *SendingWindow
	queue               *SendingQueue
	windowSize          uint32
	firstUnacknowledged uint32
	nextNumber          uint32
	remoteNextNumber    uint32
	controlWindow       uint32
	fastResend          uint32
	updated             bool
}

func NewSendingWorker(kcp *KCP) *SendingWorker {
	worker := &SendingWorker{
		kcp:              kcp,
		queue:            NewSendingQueue(effectiveConfig.GetSendingQueueSize()),
		fastResend:       2,
		remoteNextNumber: 32,
		windowSize:       effectiveConfig.GetSendingWindowSize(),
		controlWindow:    effectiveConfig.GetSendingWindowSize(),
	}
	worker.window = NewSendingWindow(effectiveConfig.GetSendingWindowSize(), effectiveConfig.GetSendingInFlightSize(), worker, worker.OnPacketLoss)
	return worker
}

func (this *SendingWorker) ProcessReceivingNext(nextNumber uint32) {
	this.Lock()
	defer this.Unlock()

	this.window.Clear(nextNumber)
	this.FindFirstUnacknowledged()
}

// @Private
func (this *SendingWorker) FindFirstUnacknowledged() {
	prevUna := this.firstUnacknowledged
	if this.window.Len() > 0 {
		this.firstUnacknowledged = this.window.First().Number
	} else {
		this.firstUnacknowledged = this.nextNumber
	}
	if this.firstUnacknowledged != prevUna {
		this.updated = true
	}
}

func (this *SendingWorker) ProcessAck(number uint32) {
	if number-this.firstUnacknowledged > this.windowSize {
		return
	}

	this.Lock()
	defer this.Unlock()
	this.window.Remove(number - this.firstUnacknowledged)
	this.FindFirstUnacknowledged()
}

func (this *SendingWorker) ProcessSegment(seg *AckSegment) {
	if this.remoteNextNumber < seg.ReceivingWindow {
		this.remoteNextNumber = seg.ReceivingWindow
	}
	this.ProcessReceivingNext(seg.ReceivingNext)
	var maxack uint32
	for i := 0; i < int(seg.Count); i++ {
		timestamp := seg.TimestampList[i]
		number := seg.NumberList[i]
		if this.kcp.current-timestamp < 10000 {
			this.kcp.update_ack(int32(this.kcp.current - timestamp))
		}
		this.ProcessAck(number)
		if maxack < number {
			maxack = number
		}
	}
	this.Lock()
	this.window.HandleFastAck(maxack)
	this.Unlock()
}

func (this *SendingWorker) Push(b []byte) int {
	nBytes := 0
	for len(b) > 0 && !this.queue.IsFull() {
		var size int
		if len(b) > int(this.kcp.mss) {
			size = int(this.kcp.mss)
		} else {
			size = len(b)
		}
		seg := &DataSegment{
			Data: alloc.NewSmallBuffer().Clear().Append(b[:size]),
		}
		this.Lock()
		this.queue.Push(seg)
		this.Unlock()
		b = b[size:]
		nBytes += size
	}
	return nBytes
}

func (this *SendingWorker) Write(seg Segment) {
	dataSeg := seg.(*DataSegment)

	dataSeg.Conv = this.kcp.conv
	dataSeg.Timestamp = this.kcp.current
	dataSeg.SendingNext = this.firstUnacknowledged
	dataSeg.Opt = 0
	if this.kcp.state == StateReadyToClose {
		dataSeg.Opt = SegmentOptionClose
	}

	this.kcp.output.Write(dataSeg)
	this.updated = false
}

func (this *SendingWorker) PingNecessary() bool {
	return this.updated
}

func (this *SendingWorker) OnPacketLoss(lossRate uint32) {
	if !effectiveConfig.Congestion || this.kcp.rx_srtt == 0 {
		return
	}

	if lossRate >= 15 {
		this.controlWindow = 3 * this.controlWindow / 4
	} else if lossRate <= 5 {
		this.controlWindow += this.controlWindow / 4
	}
	if this.controlWindow < 4 {
		this.controlWindow = 4
	}
	if this.controlWindow > 2*this.windowSize {
		this.controlWindow = 2 * this.windowSize
	}
}

func (this *SendingWorker) Flush() {
	this.Lock()
	defer this.Unlock()

	cwnd := this.firstUnacknowledged + this.windowSize
	if cwnd > this.remoteNextNumber {
		cwnd = this.remoteNextNumber
	}
	if effectiveConfig.Congestion && cwnd > this.firstUnacknowledged+this.controlWindow {
		cwnd = this.firstUnacknowledged + this.controlWindow
	}

	for !this.queue.IsEmpty() && _itimediff(this.nextNumber, cwnd) < 0 {
		seg := this.queue.Pop()
		seg.Number = this.nextNumber
		seg.timeout = this.kcp.current
		seg.ackSkipped = 0
		seg.transmit = 0
		this.window.Push(seg)
		this.nextNumber++
	}

	this.window.Flush(this.kcp.current, this.kcp.fastresend, this.kcp.rx_rto)
}

func (this *SendingWorker) CloseWrite() {
	this.Lock()
	defer this.Unlock()

	this.window.Clear(0xFFFFFFFF)
	this.queue.Clear()
}
<|endoftext|>