<commit_before>package main

import (
	"fmt"
	"strings"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/Sirupsen/tomb"

	"net"
	"time"
)

var AcceptTimeout = time.Second

// Proxy represents the proxy in its entirity with all its links. The main
// responsibility of Proxy is to accept new client and create Links between the
// client and upstream.
type Proxy struct {
	sync.Mutex

	Name     string
	Listen   string
	Upstream string

	started chan bool

	tomb  tomb.Tomb
	links []*link
}

func NewProxy() *Proxy {
	return &Proxy{
		started: make(chan bool, 1),
	}
}

func (proxy *Proxy) Start() {
	logrus.WithFields(logrus.Fields{
		"name":     proxy.Name,
		"proxy":    proxy.Listen,
		"upstream": proxy.Upstream,
	}).Info("Starting proxy")

	go proxy.server()
}

// server runs the Proxy server, accepting new clients and creating Links to
// connect them to upstreams.
func (proxy *Proxy) server() {
	ln, err := net.Listen("tcp", proxy.Listen)
	if err != nil {
		logrus.WithFields(logrus.Fields{"upstream": proxy.Upstream, "err": err}).Error("Unable to start proxy server")
		return
	}

	// This is a super hacky way to get a local address correct.
	// We want to set #Listen because if it's not supplied in the API we'll just
	// use an ephemeral port.
	tcpAddr := ln.Addr().(*net.TCPAddr)
	tcpAddrIp := strings.Trim(string(tcpAddr.IP), "\u0000")
	if tcpAddrIp == "" {
		tcpAddrIp = "localhost"
	}
	proxy.Listen = fmt.Sprintf("%s:%d", tcpAddrIp, tcpAddr.Port)

	proxy.started <- true

	for {
		// Set a deadline to not make Accept() block forever, allowing us to shut
		// down this thread.
		err = ln.(*net.TCPListener).SetDeadline(time.Now().Add(AcceptTimeout))
		if err != nil {
			logrus.WithField("name", proxy.Name).Fatal("Unable to set deadline")
		}

		// Shut down if the tomb is not empty
		select {
		case <-proxy.tomb.Dying():
			if err := ln.Close(); err != nil {
				logrus.WithFields(logrus.Fields{
					"proxy":    proxy.Listen,
					"upstream": proxy.Upstream,
					"name":     proxy.Name,
					"err":      err,
				}).Warn("Failed to shut down proxy server")
			}
			proxy.tomb.Done()
			return
		default:
		}

		client, err := ln.Accept()
		if err != nil {
			if !err.(*net.OpError).Timeout() {
				logrus.WithFields(logrus.Fields{"proxy": proxy.Listen, "err": err}).Error("Unable to accept client")
			}
			continue
		}

		logrus.WithFields(logrus.Fields{
			"name":     proxy.Name,
			"client":   client.RemoteAddr(),
			"proxy":    proxy.Listen,
			"upstream": proxy.Upstream,
		}).Info("Accepted client")

		proxy.Lock()
		link := NewLink(proxy, client)
		proxy.links = append(proxy.links, link)
		proxy.Unlock()

		if err := link.Open(); err != nil {
			logrus.WithFields(logrus.Fields{
				"name":     proxy.Name,
				"client":   client.RemoteAddr(),
				"proxy":    proxy.Listen,
				"upstream": proxy.Upstream,
			}).Error("Unable to open connection to upstream")
		}
	}
}

func (proxy *Proxy) Stop() {
	proxy.tomb.Killf("Shutting down from stop()")

	proxy.Lock()
	for _, link := range proxy.links {
		link.Close()
	}
	proxy.Unlock()

	proxy.tomb.Wait()

	logrus.WithFields(logrus.Fields{
		"name":     proxy.Name,
		"proxy":    proxy.Listen,
		"upstream": proxy.Upstream,
	}).Info("Terminated proxy")
}
<commit_msg>proxy: fix localhost parsing<commit_after>package main

import (
	"fmt"
	"sync"

	"github.com/Sirupsen/logrus"
	"github.com/Sirupsen/tomb"

	"net"
	"time"
)

var AcceptTimeout = time.Second

// Proxy represents the proxy in its entirity with all its links. The main
// responsibility of Proxy is to accept new client and create Links between the
// client and upstream.
type Proxy struct {
	sync.Mutex

	Name     string
	Listen   string
	Upstream string

	started chan bool

	tomb  tomb.Tomb
	links []*link
}

func NewProxy() *Proxy {
	return &Proxy{
		started: make(chan bool, 1),
	}
}

func (proxy *Proxy) Start() {
	go proxy.server()
}

// server runs the Proxy server, accepting new clients and creating Links to
// connect them to upstreams.
func (proxy *Proxy) server() {
	ln, err := net.Listen("tcp", proxy.Listen)
	if err != nil {
		logrus.WithFields(logrus.Fields{"upstream": proxy.Upstream, "err": err}).Error("Unable to start proxy server")
		return
	}

	// This is a super hacky way to get a local address correct.
	// We want to set #Listen because if it's not supplied in the API we'll just
	// use an ephemeral port.
	tcpAddr := ln.Addr().(*net.TCPAddr)
	tcpAddrIp := string(tcpAddr.IP)
	if net.ParseIP(string(tcpAddr.IP)) == nil {
		tcpAddrIp = "127.0.0.1"
	}
	proxy.Listen = fmt.Sprintf("%s:%d", tcpAddrIp, tcpAddr.Port)

	proxy.started <- true

	logrus.WithFields(logrus.Fields{
		"name":     proxy.Name,
		"proxy":    proxy.Listen,
		"upstream": proxy.Upstream,
	}).Info("Started proxy")

	for {
		// Set a deadline to not make Accept() block forever, allowing us to shut
		// down this thread.
		err = ln.(*net.TCPListener).SetDeadline(time.Now().Add(AcceptTimeout))
		if err != nil {
			logrus.WithField("name", proxy.Name).Fatal("Unable to set deadline")
		}

		// Shut down if the tomb is not empty
		select {
		case <-proxy.tomb.Dying():
			if err := ln.Close(); err != nil {
				logrus.WithFields(logrus.Fields{
					"proxy":    proxy.Listen,
					"upstream": proxy.Upstream,
					"name":     proxy.Name,
					"err":      err,
				}).Warn("Failed to shut down proxy server")
			}
			proxy.tomb.Done()
			return
		default:
		}

		client, err := ln.Accept()
		if err != nil {
			if !err.(*net.OpError).Timeout() {
				logrus.WithFields(logrus.Fields{"proxy": proxy.Listen, "err": err}).Error("Unable to accept client")
			}
			continue
		}

		logrus.WithFields(logrus.Fields{
			"name":     proxy.Name,
			"client":   client.RemoteAddr(),
			"proxy":    proxy.Listen,
			"upstream": proxy.Upstream,
		}).Info("Accepted client")

		proxy.Lock()
		link := NewLink(proxy, client)
		proxy.links = append(proxy.links, link)
		proxy.Unlock()

		if err := link.Open(); err != nil {
			logrus.WithFields(logrus.Fields{
				"name":     proxy.Name,
				"client":   client.RemoteAddr(),
				"proxy":    proxy.Listen,
				"upstream": proxy.Upstream,
			}).Error("Unable to open connection to upstream")
		}
	}
}

func (proxy *Proxy) Stop() {
	proxy.tomb.Killf("Shutting down from stop()")

	proxy.Lock()
	for _, link := range proxy.links {
		link.Close()
	}
	proxy.Unlock()

	proxy.tomb.Wait()

	logrus.WithFields(logrus.Fields{
		"name":     proxy.Name,
		"proxy":    proxy.Listen,
		"upstream": proxy.Upstream,
	}).Info("Terminated proxy")
}
<|endoftext|>


<commit_before>package proxy

import (
	log "github.com/Sirupsen/logrus"
	"github.com/Sirupsen/logrus/formatters/logstash"
	"net/http"
	"os"
	"regexp"
)

type ProxyConfig struct {
	Master      string
	Slaves      []string
	AwsConfig   *AWSConfig
	LogLocation string
}

func init() {
	log.SetFormatter(&logstash.LogstashFormatter{Type: "solr-proxy"})
}

type Proxy struct {
	updater *Updater
	reader  *Reader
	config  *ProxyConfig
}

func NewProxy(proxyConfig *ProxyConfig) *Proxy {
	updater := NewUpdater(proxyConfig.Master)
	reader := NewReader(proxyConfig.Slaves)

	if proxyConfig.LogLocation == "stdout" {
		log.SetOutput(os.Stdout)
	} else {
		f, err := os.OpenFile(proxyConfig.LogLocation, os.O_WRONLY|os.O_CREATE, 0755)
		if err != nil {
			panic(err)
		}
		log.SetOutput(f)
	}

	return &Proxy{
		updater: updater,
		reader:  reader,
		config:  proxyConfig,
	}
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	re := regexp.MustCompile("\\/solr\\/(.*)\\/update$")
	matches := re.FindStringSubmatch(req.URL.Path)

	req.Close = true

	if len(matches) > 0 {
		p.updater.ServeHTTP(w, req, p.config.AwsConfig, matches[1])
	} else {
		p.reader.ServeHTTP(w, req)
	}
}

func writeLog(message string, params ...string) {
	log.Printf(message, params)
}
<commit_msg>fix build<commit_after>package proxy

import (
	log "github.com/Sirupsen/logrus"
	"net/http"
	"os"
	"regexp"
)

type ProxyConfig struct {
	Master      string
	Slaves      []string
	AwsConfig   *AWSConfig
	LogLocation string
}

func init() {
	log.SetFormatter(&log.JSONFormatter{})
}

type Proxy struct {
	updater *Updater
	reader  *Reader
	config  *ProxyConfig
}

func NewProxy(proxyConfig *ProxyConfig) *Proxy {
	updater := NewUpdater(proxyConfig.Master)
	reader := NewReader(proxyConfig.Slaves)

	if proxyConfig.LogLocation == "stdout" {
		log.SetOutput(os.Stdout)
	} else {
		f, err := os.OpenFile(proxyConfig.LogLocation, os.O_WRONLY|os.O_CREATE, 0755)
		if err != nil {
			panic(err)
		}
		log.SetOutput(f)
	}

	return &Proxy{
		updater: updater,
		reader:  reader,
		config:  proxyConfig,
	}
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	re := regexp.MustCompile("\\/solr\\/(.*)\\/update$")
	matches := re.FindStringSubmatch(req.URL.Path)

	req.Close = true

	if len(matches) > 0 {
		p.updater.ServeHTTP(w, req, p.config.AwsConfig, matches[1])
	} else {
		p.reader.ServeHTTP(w, req)
	}
}

func writeLog(message string, params ...string) {
	log.Printf(message, params)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"regexp"
	"strings"
)

var xipio = regexp.MustCompile("^(.*?)\\.?\\d+\\.\\d+\\.\\d+\\.\\d+\\.xip\\.io")

// Server is the interface the wraps the Name and Port methods.
type Server interface {
	Name() string
	Port() int
}

// Servers provides a list of available servers.
type Servers interface {
	List() []Server
}

// Proxy is a ReverseProxy that takes an incoming request and
// sends it to one of the known servers based on server's name,
// after proxying the response back to the client.
type Proxy struct {
	httputil.ReverseProxy
	servers Servers
	tld     string
}

// NewProxy returns a new Proxy.
func NewProxy(tld string, s Servers) *Proxy {
	p := &Proxy{tld: tld, servers: s}
	p.Director = func(req *http.Request) {
		req.URL.Scheme = "http"
		server, found := p.resolve(req.Host)
		if found {
			// FIXME use localhost
			req.URL.Host = fmt.Sprint("127.0.0.1:", server.Port())
		} else {
			// FIXME redirect to bam server to show an error page
			log.Printf("WARN No server found for host %s\n", req.Host)
		}
	}
	return p
}

// Resolve finds a Server matching the given host.
// Return false, if any Server matches host.
func (p *Proxy) resolve(host string) (Server, bool) {
	for _, s := range p.servers.List() {
		if p.match(s, host) {
			return s, true
		}
	}
	return nil, false
}

// match checks is host matches server's name.
func (p *Proxy) match(s Server, host string) bool {
	return matchDomains(s.Name()+"."+p.tld, host) || matchXipDomain(s.Name(), host)
}

// matchDomains checks whether 'a' domain is equals to 'b' domain, or,
// 'b' domain is a subdomain of 'a' domain.
func matchDomains(a, b string) bool {
	return a == b || strings.HasSuffix(b, "."+a)
}

// matchXipDomain checks whether host is a xip domain of name.
func matchXipDomain(name, host string) bool {
	subdomain := xipio.ReplaceAllString(host, "$1")
	return matchDomains(name, subdomain)
}

type server struct {
	name string
	port int
}

func (s *server) Name() string {
	return s.name
}

func (s *server) Port() int {
	return s.port
}

func (s *server) String() string {
	return fmt.Sprintf("%s:%d", s.name, s.port)
}

func NewServer(name string, port int) Server {
	return &server{name, port}
}
<commit_msg>use localhost name instead of IPv4 address<commit_after>package main

import (
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"regexp"
	"strings"
)

var xipio = regexp.MustCompile("^(.*?)\\.?\\d+\\.\\d+\\.\\d+\\.\\d+\\.xip\\.io")

// Server is the interface the wraps the Name and Port methods.
type Server interface {
	Name() string
	Port() int
}

// Servers provides a list of available servers.
type Servers interface {
	List() []Server
}

// Proxy is a ReverseProxy that takes an incoming request and
// sends it to one of the known servers based on server's name,
// after proxying the response back to the client.
type Proxy struct {
	httputil.ReverseProxy
	servers Servers
	tld     string
}

// NewProxy returns a new Proxy.
func NewProxy(tld string, s Servers) *Proxy {
	p := &Proxy{tld: tld, servers: s}
	p.Director = func(req *http.Request) {
		req.URL.Scheme = "http"
		server, found := p.resolve(req.Host)
		if found {
			req.URL.Host = fmt.Sprint("localhost:", server.Port())
		} else {
			// FIXME redirect to bam server to show an error page
			log.Printf("WARN No server found for host %s\n", req.Host)
		}
	}
	return p
}

// Resolve finds a Server matching the given host.
// Return false, if any Server matches host.
func (p *Proxy) resolve(host string) (Server, bool) {
	for _, s := range p.servers.List() {
		if p.match(s, host) {
			return s, true
		}
	}
	return nil, false
}

// match checks is host matches server's name.
func (p *Proxy) match(s Server, host string) bool {
	return matchDomains(s.Name()+"."+p.tld, host) || matchXipDomain(s.Name(), host)
}

// matchDomains checks whether 'a' domain is equals to 'b' domain, or,
// 'b' domain is a subdomain of 'a' domain.
func matchDomains(a, b string) bool {
	return a == b || strings.HasSuffix(b, "."+a)
}

// matchXipDomain checks whether host is a xip domain of name.
func matchXipDomain(name, host string) bool {
	subdomain := xipio.ReplaceAllString(host, "$1")
	return matchDomains(name, subdomain)
}

type server struct {
	name string
	port int
}

func (s *server) Name() string {
	return s.name
}

func (s *server) Port() int {
	return s.port
}

func (s *server) String() string {
	return fmt.Sprintf("%s:%d", s.name, s.port)
}

func NewServer(name string, port int) Server {
	return &server{name, port}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"

	"github.com/cloudflare/unsee/internal/alertmanager"
	"github.com/cloudflare/unsee/internal/config"
	"github.com/gin-gonic/gin"

	log "github.com/sirupsen/logrus"
)

func proxyPathPrefix(name string) string {
	return fmt.Sprintf("%sproxy/alertmanager/%s", config.Config.Listen.Prefix, name)
}

func proxyPath(name, path string) string {
	return fmt.Sprintf("%s%s", proxyPathPrefix(name), path)
}

// NewAlertmanagerProxy creates a proxy instance for given alertmanager instance
func NewAlertmanagerProxy(alertmanager *alertmanager.Alertmanager) (*httputil.ReverseProxy, error) {
	upstreamURL, err := url.Parse(alertmanager.URI)
	if err != nil {
		return nil, err
	}
	proxy := httputil.ReverseProxy{
		Director: func(req *http.Request) {
			req.URL.Scheme = upstreamURL.Scheme
			req.URL.Host = upstreamURL.Host
			// drop Accept-Encoding header so we always get uncompressed reponses from
			// upstream, there's a gzip middleware that's global so we don't want it
			// to gzip twice
			req.Header.Del("Accept-Encoding")
			log.Debugf("[%s] Proxy request for %s", alertmanager.Name, req.URL.Path)
		},
		Transport: alertmanager.HTTPTransport,
		ModifyResponse: func(resp *http.Response) error {
			// drop Content-Length header from upstream responses, gzip middleware
			// will compress those and that could cause a mismatch
			resp.Header.Del("Content-Length")
			return nil
		},
	}
	return &proxy, nil
}

func setupRouterProxyHandlers(router *gin.Engine, alertmanager *alertmanager.Alertmanager) error {
	proxy, err := NewAlertmanagerProxy(alertmanager)
	if err != nil {
		return err
	}
	router.POST(
		proxyPath(alertmanager.Name, "/api/v1/silences"),
		gin.WrapH(http.StripPrefix(proxyPathPrefix(alertmanager.Name), proxy)))
	router.DELETE(
		proxyPath(alertmanager.Name, "/api/v1/silence/*id"),
		gin.WrapH(http.StripPrefix(proxyPathPrefix(alertmanager.Name), proxy)))
	return nil
}
<commit_msg>Remove Host header attribute from proxied request (#252)<commit_after>package main

import (
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"

	"github.com/cloudflare/unsee/internal/alertmanager"
	"github.com/cloudflare/unsee/internal/config"
	"github.com/gin-gonic/gin"

	log "github.com/sirupsen/logrus"
)

func proxyPathPrefix(name string) string {
	return fmt.Sprintf("%sproxy/alertmanager/%s", config.Config.Listen.Prefix, name)
}

func proxyPath(name, path string) string {
	return fmt.Sprintf("%s%s", proxyPathPrefix(name), path)
}

// NewAlertmanagerProxy creates a proxy instance for given alertmanager instance
func NewAlertmanagerProxy(alertmanager *alertmanager.Alertmanager) (*httputil.ReverseProxy, error) {
	upstreamURL, err := url.Parse(alertmanager.URI)
	if err != nil {
		return nil, err
	}
	proxy := httputil.ReverseProxy{
		Director: func(req *http.Request) {
			req.URL.Scheme = upstreamURL.Scheme
			req.URL.Host = upstreamURL.Host
			// drop Accept-Encoding header so we always get uncompressed reponses from
			// upstream, there's a gzip middleware that's global so we don't want it
			// to gzip twice
			req.Header.Del("Accept-Encoding")

			// set hostname of proxied target
			req.Host = upstreamURL.Host

			log.Debugf("[%s] Proxy request for %s", alertmanager.Name, req.URL.Path)
		},
		Transport: alertmanager.HTTPTransport,
		ModifyResponse: func(resp *http.Response) error {
			// drop Content-Length header from upstream responses, gzip middleware
			// will compress those and that could cause a mismatch
			resp.Header.Del("Content-Length")
			return nil
		},
	}
	return &proxy, nil
}

func setupRouterProxyHandlers(router *gin.Engine, alertmanager *alertmanager.Alertmanager) error {
	proxy, err := NewAlertmanagerProxy(alertmanager)
	if err != nil {
		return err
	}
	router.POST(
		proxyPath(alertmanager.Name, "/api/v1/silences"),
		gin.WrapH(http.StripPrefix(proxyPathPrefix(alertmanager.Name), proxy)))
	router.DELETE(
		proxyPath(alertmanager.Name, "/api/v1/silence/*id"),
		gin.WrapH(http.StripPrefix(proxyPathPrefix(alertmanager.Name), proxy)))
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package pe implements access to PE (Microsoft Windows Portable Executable) files.
package pe

import (
	"debug/dwarf"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"strconv"
)

// A File represents an open PE file.
type File struct {
	FileHeader
	Sections []*Section

	closer io.Closer
}

type SectionHeader struct {
	Name                 string
	VirtualSize          uint32
	VirtualAddress       uint32
	Size                 uint32
	Offset               uint32
	PointerToRelocations uint32
	PointerToLineNumbers uint32
	NumberOfRelocations  uint16
	NumberOfLineNumbers  uint16
	Characteristics      uint32
}

type Section struct {
	SectionHeader

	// Embed ReaderAt for ReadAt method.
	// Do not embed SectionReader directly
	// to avoid having Read and Seek.
	// If a client wants Read and Seek it must use
	// Open() to avoid fighting over the seek offset
	// with other clients.
	io.ReaderAt
	sr *io.SectionReader
}

type ImportDirectory struct {
	OriginalFirstThunk uint32
	TimeDateStamp      uint32
	ForwarderChain     uint32
	Name               uint32
	FirstThunk         uint32

	dll string
	rva []uint32
}

// Data reads and returns the contents of the PE section.
func (s *Section) Data() ([]byte, os.Error) {
	dat := make([]byte, s.sr.Size())
	n, err := s.sr.ReadAt(dat, 0)
	return dat[0:n], err
}

// Open returns a new ReadSeeker reading the PE section.
func (s *Section) Open() io.ReadSeeker { return io.NewSectionReader(s.sr, 0, 1<<63-1) }

type FormatError struct {
	off int64
	msg string
	val interface{}
}

func (e *FormatError) String() string {
	msg := e.msg
	if e.val != nil {
		msg += fmt.Sprintf(" '%v'", e.val)
	}
	msg += fmt.Sprintf(" in record at byte %#x", e.off)
	return msg
}

// Open opens the named file using os.Open and prepares it for use as a PE binary.
func Open(name string) (*File, os.Error) {
	f, err := os.Open(name, os.O_RDONLY, 0)
	if err != nil {
		return nil, err
	}
	ff, err := NewFile(f)
	if err != nil {
		f.Close()
		return nil, err
	}
	ff.closer = f
	return ff, nil
}

// Close closes the File.
// If the File was created using NewFile directly instead of Open,
// Close has no effect.
func (f *File) Close() os.Error {
	var err os.Error
	if f.closer != nil {
		err = f.closer.Close()
		f.closer = nil
	}
	return err
}

// NewFile creates a new File for acecssing a PE binary in an underlying reader.
func NewFile(r io.ReaderAt) (*File, os.Error) {
	f := new(File)
	sr := io.NewSectionReader(r, 0, 1<<63-1)

	var dosheader [96]byte
	if _, err := r.ReadAt(dosheader[0:], 0); err != nil {
		return nil, err
	}
	var base int64
	if dosheader[0] == 'M' && dosheader[1] == 'Z' {
		var sign [4]byte
		r.ReadAt(sign[0:], int64(dosheader[0x3c]))
		if !(sign[0] == 'P' && sign[1] == 'E' && sign[2] == 0 && sign[3] == 0) {
			return nil, os.NewError("Invalid PE File Format.")
		}
		base = int64(dosheader[0x3c]) + 4
	} else {
		base = int64(0)
	}
	sr.Seek(base, 0)
	if err := binary.Read(sr, binary.LittleEndian, &f.FileHeader); err != nil {
		return nil, err
	}
	if f.FileHeader.Machine != IMAGE_FILE_MACHINE_UNKNOWN && f.FileHeader.Machine != IMAGE_FILE_MACHINE_AMD64 && f.FileHeader.Machine != IMAGE_FILE_MACHINE_I386 {
		return nil, os.NewError("Invalid PE File Format.")
	}
	// get symbol string table
	sr.Seek(int64(f.FileHeader.PointerToSymbolTable+18*f.FileHeader.NumberOfSymbols), 0)
	var l uint32
	if err := binary.Read(sr, binary.LittleEndian, &l); err != nil {
		return nil, err
	}
	ss := make([]byte, l)
	if _, err := r.ReadAt(ss, int64(f.FileHeader.PointerToSymbolTable+18*f.FileHeader.NumberOfSymbols)); err != nil {
		return nil, err
	}
	sr.Seek(base, 0)
	binary.Read(sr, binary.LittleEndian, &f.FileHeader)
	sr.Seek(int64(f.FileHeader.SizeOfOptionalHeader), 1) //Skip OptionalHeader
	f.Sections = make([]*Section, f.FileHeader.NumberOfSections)
	for i := 0; i < int(f.FileHeader.NumberOfSections); i++ {
		sh := new(SectionHeader32)
		if err := binary.Read(sr, binary.LittleEndian, sh); err != nil {
			return nil, err
		}
		var name string
		if sh.Name[0] == '\x2F' {
			si, _ := strconv.Atoi(cstring(sh.Name[1:]))
			name, _ = getString(ss, si)
		} else {
			name = cstring(sh.Name[0:])
		}
		s := new(Section)
		s.SectionHeader = SectionHeader{
			Name:                 name,
			VirtualSize:          uint32(sh.VirtualSize),
			VirtualAddress:       uint32(sh.VirtualAddress),
			Size:                 uint32(sh.SizeOfRawData),
			Offset:               uint32(sh.PointerToRawData),
			PointerToRelocations: uint32(sh.PointerToRelocations),
			PointerToLineNumbers: uint32(sh.PointerToLineNumbers),
			NumberOfRelocations:  uint16(sh.NumberOfRelocations),
			NumberOfLineNumbers:  uint16(sh.NumberOfLineNumbers),
			Characteristics:      uint32(sh.Characteristics),
		}
		s.sr = io.NewSectionReader(r, int64(s.SectionHeader.Offset), int64(s.SectionHeader.Size))
		s.ReaderAt = s.sr
		f.Sections[i] = s
	}
	return f, nil
}

func cstring(b []byte) string {
	var i int
	for i = 0; i < len(b) && b[i] != 0; i++ {
	}
	return string(b[0:i])
}

// getString extracts a string from symbol string table.
func getString(section []byte, start int) (string, bool) {
	if start < 0 || start >= len(section) {
		return "", false
	}

	for end := start; end < len(section); end++ {
		if section[end] == 0 {
			return string(section[start:end]), true
		}
	}
	return "", false
}

// Section returns the first section with the given name, or nil if no such
// section exists.
func (f *File) Section(name string) *Section {
	for _, s := range f.Sections {
		if s.Name == name {
			return s
		}
	}
	return nil
}

func (f *File) DWARF() (*dwarf.Data, os.Error) {
	// There are many other DWARF sections, but these
	// are the required ones, and the debug/dwarf package
	// does not use the others, so don't bother loading them.
	var names = [...]string{"abbrev", "info", "str"}
	var dat [len(names)][]byte
	for i, name := range names {
		name = ".debug_" + name
		s := f.Section(name)
		if s == nil {
			continue
		}
		b, err := s.Data()
		if err != nil && uint32(len(b)) < s.Size {
			return nil, err
		}
		dat[i] = b
	}

	abbrev, info, str := dat[0], dat[1], dat[2]
	return dwarf.New(abbrev, nil, nil, info, nil, nil, nil, str)
}

// ImportedSymbols returns the names of all symbols
// referred to by the binary f that are expected to be
// satisfied by other libraries at dynamic load time.
// It does not return weak symbols.
func (f *File) ImportedSymbols() ([]string, os.Error) {
	ds := f.Section(".idata")
	if ds == nil {
		// not dynamic, so no libraries
		return nil, nil
	}
	d, err := ds.Data()
	if err != nil {
		return nil, err
	}
	var ida []ImportDirectory
	for len(d) > 0 {
		var dt ImportDirectory
		dt.OriginalFirstThunk = binary.LittleEndian.Uint32(d[0:4])
		dt.Name = binary.LittleEndian.Uint32(d[12:16])
		dt.FirstThunk = binary.LittleEndian.Uint32(d[16:20])
		d = d[20:]
		if dt.OriginalFirstThunk == 0 {
			break
		}
		ida = append(ida, dt)
	}
	for i, _ := range ida {
		for len(d) > 0 {
			va := binary.LittleEndian.Uint32(d[0:4])
			d = d[4:]
			if va == 0 {
				break
			}
			ida[i].rva = append(ida[i].rva, va)
		}
	}
	for _, _ = range ida {
		for len(d) > 0 {
			va := binary.LittleEndian.Uint32(d[0:4])
			d = d[4:]
			if va == 0 {
				break
			}
		}
	}
	names, _ := ds.Data()
	var all []string
	for _, dt := range ida {
		dt.dll, _ = getString(names, int(dt.Name-ds.VirtualAddress))
		for _, va := range dt.rva {
			fn, _ := getString(names, int(va-ds.VirtualAddress+2))
			all = append(all, fn+":"+dt.dll)
		}
	}

	return all, nil
}

// ImportedLibraries returns the names of all libraries
// referred to by the binary f that are expected to be
// linked with the binary at dynamic link time.
func (f *File) ImportedLibraries() ([]string, os.Error) {
	// TODO
	// cgo -dynimport don't use this for windows PE, so just return.
	return nil, nil
}
<commit_msg>debug/pe: ImportedSymbols fixes<commit_after>// Copyright 2009 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package pe implements access to PE (Microsoft Windows Portable Executable) files.
package pe

import (
	"debug/dwarf"
	"encoding/binary"
	"fmt"
	"io"
	"os"
	"strconv"
)

// A File represents an open PE file.
type File struct {
	FileHeader
	Sections []*Section

	closer io.Closer
}

type SectionHeader struct {
	Name                 string
	VirtualSize          uint32
	VirtualAddress       uint32
	Size                 uint32
	Offset               uint32
	PointerToRelocations uint32
	PointerToLineNumbers uint32
	NumberOfRelocations  uint16
	NumberOfLineNumbers  uint16
	Characteristics      uint32
}

type Section struct {
	SectionHeader

	// Embed ReaderAt for ReadAt method.
	// Do not embed SectionReader directly
	// to avoid having Read and Seek.
	// If a client wants Read and Seek it must use
	// Open() to avoid fighting over the seek offset
	// with other clients.
	io.ReaderAt
	sr *io.SectionReader
}

type ImportDirectory struct {
	OriginalFirstThunk uint32
	TimeDateStamp      uint32
	ForwarderChain     uint32
	Name               uint32
	FirstThunk         uint32

	dll string
}

// Data reads and returns the contents of the PE section.
func (s *Section) Data() ([]byte, os.Error) {
	dat := make([]byte, s.sr.Size())
	n, err := s.sr.ReadAt(dat, 0)
	return dat[0:n], err
}

// Open returns a new ReadSeeker reading the PE section.
func (s *Section) Open() io.ReadSeeker { return io.NewSectionReader(s.sr, 0, 1<<63-1) }

type FormatError struct {
	off int64
	msg string
	val interface{}
}

func (e *FormatError) String() string {
	msg := e.msg
	if e.val != nil {
		msg += fmt.Sprintf(" '%v'", e.val)
	}
	msg += fmt.Sprintf(" in record at byte %#x", e.off)
	return msg
}

// Open opens the named file using os.Open and prepares it for use as a PE binary.
func Open(name string) (*File, os.Error) {
	f, err := os.Open(name, os.O_RDONLY, 0)
	if err != nil {
		return nil, err
	}
	ff, err := NewFile(f)
	if err != nil {
		f.Close()
		return nil, err
	}
	ff.closer = f
	return ff, nil
}

// Close closes the File.
// If the File was created using NewFile directly instead of Open,
// Close has no effect.
func (f *File) Close() os.Error {
	var err os.Error
	if f.closer != nil {
		err = f.closer.Close()
		f.closer = nil
	}
	return err
}

// NewFile creates a new File for acecssing a PE binary in an underlying reader.
func NewFile(r io.ReaderAt) (*File, os.Error) {
	f := new(File)
	sr := io.NewSectionReader(r, 0, 1<<63-1)

	var dosheader [96]byte
	if _, err := r.ReadAt(dosheader[0:], 0); err != nil {
		return nil, err
	}
	var base int64
	if dosheader[0] == 'M' && dosheader[1] == 'Z' {
		var sign [4]byte
		r.ReadAt(sign[0:], int64(dosheader[0x3c]))
		if !(sign[0] == 'P' && sign[1] == 'E' && sign[2] == 0 && sign[3] == 0) {
			return nil, os.NewError("Invalid PE File Format.")
		}
		base = int64(dosheader[0x3c]) + 4
	} else {
		base = int64(0)
	}
	sr.Seek(base, 0)
	if err := binary.Read(sr, binary.LittleEndian, &f.FileHeader); err != nil {
		return nil, err
	}
	if f.FileHeader.Machine != IMAGE_FILE_MACHINE_UNKNOWN && f.FileHeader.Machine != IMAGE_FILE_MACHINE_AMD64 && f.FileHeader.Machine != IMAGE_FILE_MACHINE_I386 {
		return nil, os.NewError("Invalid PE File Format.")
	}
	// get symbol string table
	sr.Seek(int64(f.FileHeader.PointerToSymbolTable+18*f.FileHeader.NumberOfSymbols), 0)
	var l uint32
	if err := binary.Read(sr, binary.LittleEndian, &l); err != nil {
		return nil, err
	}
	ss := make([]byte, l)
	if _, err := r.ReadAt(ss, int64(f.FileHeader.PointerToSymbolTable+18*f.FileHeader.NumberOfSymbols)); err != nil {
		return nil, err
	}
	sr.Seek(base, 0)
	binary.Read(sr, binary.LittleEndian, &f.FileHeader)
	sr.Seek(int64(f.FileHeader.SizeOfOptionalHeader), 1) //Skip OptionalHeader
	f.Sections = make([]*Section, f.FileHeader.NumberOfSections)
	for i := 0; i < int(f.FileHeader.NumberOfSections); i++ {
		sh := new(SectionHeader32)
		if err := binary.Read(sr, binary.LittleEndian, sh); err != nil {
			return nil, err
		}
		var name string
		if sh.Name[0] == '\x2F' {
			si, _ := strconv.Atoi(cstring(sh.Name[1:]))
			name, _ = getString(ss, si)
		} else {
			name = cstring(sh.Name[0:])
		}
		s := new(Section)
		s.SectionHeader = SectionHeader{
			Name:                 name,
			VirtualSize:          uint32(sh.VirtualSize),
			VirtualAddress:       uint32(sh.VirtualAddress),
			Size:                 uint32(sh.SizeOfRawData),
			Offset:               uint32(sh.PointerToRawData),
			PointerToRelocations: uint32(sh.PointerToRelocations),
			PointerToLineNumbers: uint32(sh.PointerToLineNumbers),
			NumberOfRelocations:  uint16(sh.NumberOfRelocations),
			NumberOfLineNumbers:  uint16(sh.NumberOfLineNumbers),
			Characteristics:      uint32(sh.Characteristics),
		}
		s.sr = io.NewSectionReader(r, int64(s.SectionHeader.Offset), int64(s.SectionHeader.Size))
		s.ReaderAt = s.sr
		f.Sections[i] = s
	}
	return f, nil
}

func cstring(b []byte) string {
	var i int
	for i = 0; i < len(b) && b[i] != 0; i++ {
	}
	return string(b[0:i])
}

// getString extracts a string from symbol string table.
func getString(section []byte, start int) (string, bool) {
	if start < 0 || start >= len(section) {
		return "", false
	}

	for end := start; end < len(section); end++ {
		if section[end] == 0 {
			return string(section[start:end]), true
		}
	}
	return "", false
}

// Section returns the first section with the given name, or nil if no such
// section exists.
func (f *File) Section(name string) *Section {
	for _, s := range f.Sections {
		if s.Name == name {
			return s
		}
	}
	return nil
}

func (f *File) DWARF() (*dwarf.Data, os.Error) {
	// There are many other DWARF sections, but these
	// are the required ones, and the debug/dwarf package
	// does not use the others, so don't bother loading them.
	var names = [...]string{"abbrev", "info", "str"}
	var dat [len(names)][]byte
	for i, name := range names {
		name = ".debug_" + name
		s := f.Section(name)
		if s == nil {
			continue
		}
		b, err := s.Data()
		if err != nil && uint32(len(b)) < s.Size {
			return nil, err
		}
		dat[i] = b
	}

	abbrev, info, str := dat[0], dat[1], dat[2]
	return dwarf.New(abbrev, nil, nil, info, nil, nil, nil, str)
}

// ImportedSymbols returns the names of all symbols
// referred to by the binary f that are expected to be
// satisfied by other libraries at dynamic load time.
// It does not return weak symbols.
func (f *File) ImportedSymbols() ([]string, os.Error) {
	ds := f.Section(".idata")
	if ds == nil {
		// not dynamic, so no libraries
		return nil, nil
	}
	d, err := ds.Data()
	if err != nil {
		return nil, err
	}
	var ida []ImportDirectory
	for len(d) > 0 {
		var dt ImportDirectory
		dt.OriginalFirstThunk = binary.LittleEndian.Uint32(d[0:4])
		dt.Name = binary.LittleEndian.Uint32(d[12:16])
		dt.FirstThunk = binary.LittleEndian.Uint32(d[16:20])
		d = d[20:]
		if dt.OriginalFirstThunk == 0 {
			break
		}
		ida = append(ida, dt)
	}
	names, _ := ds.Data()
	var all []string
	for _, dt := range ida {
		dt.dll, _ = getString(names, int(dt.Name-ds.VirtualAddress))
		d, _ = ds.Data()
		// seek to OriginalFirstThunk
		d = d[dt.OriginalFirstThunk-ds.VirtualAddress:]
		for len(d) > 0 {
			va := binary.LittleEndian.Uint32(d[0:4])
			d = d[4:]
			if va == 0 {
				break
			}
			if va&0x80000000 > 0 { // is Ordinal
				// TODO add dynimport ordinal support.
				//ord := va&0x0000FFFF
			} else {
				fn, _ := getString(names, int(va-ds.VirtualAddress+2))
				all = append(all, fn+":"+dt.dll)
			}
		}
	}

	return all, nil
}

// ImportedLibraries returns the names of all libraries
// referred to by the binary f that are expected to be
// linked with the binary at dynamic link time.
func (f *File) ImportedLibraries() ([]string, os.Error) {
	// TODO
	// cgo -dynimport don't use this for windows PE, so just return.
	return nil, nil
}
<|endoftext|>


<commit_before>package periodic

import (
    "log"
    "net"
    "time"
    "sync"
    "strings"
    "container/list"
    "container/heap"
    "github.com/Lupino/periodic/driver"
    "github.com/Lupino/periodic/protocol"
)

type Sched struct {
    jobTimer   *time.Timer
    grabQueue  *GrabQueue
    procQueue  *list.List
    revertPQ   PriorityQueue
    revTimer   *time.Timer
    entryPoint string
    JobLocker  *sync.Mutex
    stats      map[string]*FuncStat
    FuncLocker *sync.Mutex
    driver     driver.StoreDriver
    jobPQ      map[string]*PriorityQueue
    PQLocker   *sync.Mutex
    timeout    time.Duration
    alive      bool
}

func NewSched(entryPoint string, driver driver.StoreDriver, timeout time.Duration) *Sched {
    sched := new(Sched)
    sched.jobTimer = time.NewTimer(1 * time.Hour)
    sched.revTimer = time.NewTimer(1 * time.Hour)
    sched.grabQueue = NewGrabQueue()
    sched.procQueue = list.New()
    sched.revertPQ = make(PriorityQueue, 0)
    heap.Init(&sched.revertPQ)
    sched.entryPoint = entryPoint
    sched.JobLocker = new(sync.Mutex)
    sched.PQLocker = new(sync.Mutex)
    sched.FuncLocker = new(sync.Mutex)
    sched.stats = make(map[string]*FuncStat)
    sched.driver = driver
    sched.jobPQ = make(map[string]*PriorityQueue)
    sched.timeout = timeout
    sched.alive = true
    return sched
}

func (sched *Sched) Serve() {
    parts := strings.SplitN(sched.entryPoint, "://", 2)
    if parts[0] == "unix" {
        sockCheck(parts[1])
    }
    sched.loadJobQueue()
    go sched.handleJobPQ()
    go sched.handleRevertPQ()
    listen, err := net.Listen(parts[0], parts[1])
    if err != nil {
        log.Fatal(err)
    }
    defer listen.Close()
    log.Printf("Periodic task system started on %s\n", sched.entryPoint)
    for {
        if !sched.alive {
            break
        }
        conn, err := listen.Accept()
        if err != nil {
            log.Fatal(err)
        }
        if sched.timeout > 0 {
            conn.SetDeadline(time.Now().Add(sched.timeout * time.Second))
        }
        sched.HandleConnection(conn)
    }
}

func (sched *Sched) NotifyJobTimer() {
    sched.jobTimer.Reset(time.Millisecond)
}

func (sched *Sched) NotifyRevertTimer() {
    sched.revTimer.Reset(time.Millisecond)
}

func (sched *Sched) HandleConnection(conn net.Conn) {
    c := Conn{Conn: conn}
    payload, err := c.Receive()
    if err != nil {
        return
    }
    switch protocol.ClientType(payload[0]) {
    case protocol.TYPE_CLIENT:
        client := NewClient(sched, c)
        go client.Handle()
        break
    case protocol.TYPE_WORKER:
        worker := NewWorker(sched, c)
        go worker.Handle()
        break
    default:
        log.Printf("Unsupport client %d\n", payload[0])
        c.Close()
        break
    }
}

func (sched *Sched) Done(jobId int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    removeListJob(sched.procQueue, jobId)
    job, err := sched.driver.Get(jobId)
    if err == nil {
        sched.driver.Delete(jobId)
        sched.DecrStatJob(job)
        sched.DecrStatProc(job)
        sched.removeRevertPQ(job)
    }
    return
}

func (sched *Sched) isDoJob(job driver.Job) bool {
    for e := sched.procQueue.Front(); e != nil; e = e.Next() {
        chk := e.Value.(driver.Job)
        if chk.Id == job.Id {
            return true
        }
    }
    return false
}

func (sched *Sched) SubmitJob(grabItem GrabItem, job driver.Job) bool {
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    if job.Name == "" {
        sched.driver.Delete(job.Id)
        return true
    }
    if sched.isDoJob(job) {
        return true
    }
    if !grabItem.w.alive {
        return false
    }
    if err := grabItem.w.HandleDo(grabItem.msgId, job); err != nil {
        grabItem.w.alive = false
        return false
    }
    now := time.Now()
    current := int64(now.Unix())
    job.Status = driver.JOB_STATUS_PROC
    job.RunAt = current
    sched.driver.Save(&job)
    sched.IncrStatProc(job)
    sched.pushRevertPQ(job)
    sched.NotifyRevertTimer()
    sched.procQueue.PushBack(job)
    sched.grabQueue.Remove(grabItem)
    return true
}

func (sched *Sched) lessItem() (lessItem *Item) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    maybeItem := make(map[string]*Item)
    for Func, stat := range sched.stats {
        if stat.Worker == 0 {
            continue
        }
        pq, ok := sched.jobPQ[Func]
        if !ok || pq.Len() == 0 {
            continue
        }

        item := heap.Pop(pq).(*Item)

        maybeItem[Func] = item

    }

    if len(maybeItem) == 0 {
        return nil
    }

    var lessFunc string

    for Func, item := range maybeItem {
        if lessItem == nil {
            lessItem = item
            lessFunc = Func
            continue
        }
        if lessItem.priority > item.priority {
            lessItem = item
            lessFunc = Func
        }
    }

    for Func, item := range maybeItem {
        if Func == lessFunc {
            continue
        }
        pq := sched.jobPQ[Func]
        heap.Push(pq, item)
    }
    return
}

func (sched *Sched) handleJobPQ() {
    var current time.Time
    var timestamp int64
    for {
        if !sched.alive {
            break
        }
        if sched.grabQueue.Len() == 0 {
            sched.jobTimer.Reset(time.Minute)
            current =<-sched.jobTimer.C
            continue
        }

        lessItem := sched.lessItem()

        if lessItem == nil {
            sched.jobTimer.Reset(time.Minute)
            current =<-sched.jobTimer.C
            continue
        }

        schedJob, err := sched.driver.Get(lessItem.value)

        if err != nil {
            log.Printf("Error: Get job: %d %v\n", lessItem.value, err)
            continue
        }

        timestamp = int64(time.Now().Unix())

        if schedJob.SchedAt > timestamp {
            sched.jobTimer.Reset(time.Second * time.Duration(schedJob.SchedAt - timestamp))
            current =<-sched.jobTimer.C
            timestamp = int64(current.Unix())
            if schedJob.SchedAt > timestamp {
                sched.pushJobPQ(schedJob)
                continue
            }
        }

        grabItem, err := sched.grabQueue.Get(schedJob.Func)
        if err == nil {
            if !sched.SubmitJob(grabItem, schedJob) {
                sched.pushJobPQ(schedJob)
            }
        } else {
            sched.pushJobPQ(schedJob)
        }
    }
}

func (sched *Sched) handleRevertPQ() {
    var current time.Time
    var timestamp int64
    for {
        if !sched.alive {
            break
        }
        if sched.revertPQ.Len() == 0 {
            sched.revTimer.Reset(time.Minute)
            current =<-sched.revTimer.C
            continue
        }

        sched.PQLocker.Lock()
        item := heap.Pop(&sched.revertPQ).(*Item)
        sched.PQLocker.Unlock()

        if item == nil {
            sched.revTimer.Reset(time.Minute)
            current =<-sched.revTimer.C
            continue
        }

        revertJob, err := sched.driver.Get(item.value)

        if err != nil {
            log.Printf("Error: Get job: %d %v\n", item.value, err)
            continue
        }

        timestamp = int64(time.Now().Unix())

        if item.priority > timestamp {
            sched.revTimer.Reset(time.Second * time.Duration(item.priority - timestamp))
            current =<-sched.revTimer.C
            timestamp = int64(current.Unix())
            if item.priority > timestamp {
                sched.pushRevertPQ(revertJob)
                continue
            }
        }

        sched.DecrStatProc(revertJob)
        revertJob.Status = driver.JOB_STATUS_READY
        sched.driver.Save(&revertJob)
        sched.pushJobPQ(revertJob)
        removeListJob(sched.procQueue, revertJob.Id)
    }
}

func (sched *Sched) Fail(jobId int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    removeListJob(sched.procQueue, jobId)
    job, _ := sched.driver.Get(jobId)
    sched.DecrStatProc(job)
    sched.removeRevertPQ(job)
    job.Status = driver.JOB_STATUS_READY
    sched.driver.Save(&job)
    sched.pushJobPQ(job)
    return
}

func (sched *Sched) getFuncStat(Func string) *FuncStat {
    defer sched.FuncLocker.Unlock()
    sched.FuncLocker.Lock()
    stat, ok := sched.stats[Func]
    if !ok {
        stat = new(FuncStat)
        sched.stats[Func] = stat
    }
    return stat
}

func (sched *Sched) IncrStatFunc(Func string) {
    stat := sched.getFuncStat(Func)
    stat.Worker.Incr()
}

func (sched *Sched) DecrStatFunc(Func string) {
    stat := sched.getFuncStat(Func)
    stat.Worker.Decr()
}

func (sched *Sched) IncrStatJob(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    stat.Job.Incr()
}

func (sched *Sched) DecrStatJob(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    stat.Job.Decr()
}

func (sched *Sched) IncrStatProc(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    if job.Status == driver.JOB_STATUS_PROC {
        stat.Processing.Incr()
    }
}

func (sched *Sched) DecrStatProc(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    if job.Status == driver.JOB_STATUS_PROC {
        stat.Processing.Decr()
    }
}

func (sched *Sched) SchedLater(jobId int64, delay int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    removeListJob(sched.procQueue, jobId)
    job, _ := sched.driver.Get(jobId)
    sched.DecrStatProc(job)
    sched.removeRevertPQ(job)
    job.Status = driver.JOB_STATUS_READY
    var now = time.Now()
    job.SchedAt = int64(now.Unix()) + delay
    sched.driver.Save(&job)
    sched.pushJobPQ(job)
    return
}

func (sched *Sched) pushJobPQ(job driver.Job) bool {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_READY {
        pq, ok := sched.jobPQ[job.Func]
        if !ok {
            pq1 := make(PriorityQueue, 0)
            pq = &pq1
            sched.jobPQ[job.Func] = pq
            heap.Init(pq)
        }
        item := &Item{
            value: job.Id,
            priority: job.SchedAt,
        }
        heap.Push(pq, item)
        return true
    }
    return false
}

func (sched *Sched) pushRevertPQ(job driver.Job) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_PROC && job.Timeout > 0 {
        runAt := job.RunAt
        if runAt == 0 {
            runAt = job.SchedAt
        }
        item := &Item{
            value: job.Id,
            priority: runAt + job.Timeout,
        }
        heap.Push(&sched.revertPQ, item)
    }
}

func (sched *Sched) removeRevertPQ(job driver.Job) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_PROC && job.Timeout > 0 {
        for _, item := range sched.revertPQ {
            if item.value == job.Id {
                heap.Remove(&sched.revertPQ, item.index)
                break
            }
        }
    }
}

func (sched *Sched) loadJobQueue() {
    updateQueue := make([]driver.Job, 0)
    removeQueue := make([]driver.Job, 0)
    var now = time.Now()
    current := int64(now.Unix())

    iter := sched.driver.NewIterator(nil)
    for {
        if !iter.Next() {
            break
        }
        job := iter.Value()
        if job.Name == "" {
            removeQueue = append(removeQueue, job)
            continue
        }
        sched.IncrStatJob(job)
        sched.pushJobPQ(job)
        runAt := job.RunAt
        if runAt < job.SchedAt {
            runAt = job.SchedAt
        }
        if runAt + job.Timeout < current {
            updateQueue = append(updateQueue, job)
        } else {
            sched.procQueue.PushBack(job)
            sched.IncrStatProc(job)
            sched.pushRevertPQ(job)
        }
    }

    iter.Close()

    for _, job := range updateQueue {
        job.Status = driver.JOB_STATUS_READY
        sched.driver.Save(&job)
    }

    for _, job := range removeQueue {
        sched.driver.Delete(job.Id)
    }
}

func (sched *Sched) Close() {
    sched.alive = false
    sched.driver.Close()
    log.Printf("Periodic task system shutdown\n")
}
<commit_msg>update error msg<commit_after>package periodic

import (
    "log"
    "net"
    "time"
    "sync"
    "strings"
    "container/list"
    "container/heap"
    "github.com/Lupino/periodic/driver"
    "github.com/Lupino/periodic/protocol"
)

type Sched struct {
    jobTimer   *time.Timer
    grabQueue  *GrabQueue
    procQueue  *list.List
    revertPQ   PriorityQueue
    revTimer   *time.Timer
    entryPoint string
    JobLocker  *sync.Mutex
    stats      map[string]*FuncStat
    FuncLocker *sync.Mutex
    driver     driver.StoreDriver
    jobPQ      map[string]*PriorityQueue
    PQLocker   *sync.Mutex
    timeout    time.Duration
    alive      bool
}

func NewSched(entryPoint string, driver driver.StoreDriver, timeout time.Duration) *Sched {
    sched := new(Sched)
    sched.jobTimer = time.NewTimer(1 * time.Hour)
    sched.revTimer = time.NewTimer(1 * time.Hour)
    sched.grabQueue = NewGrabQueue()
    sched.procQueue = list.New()
    sched.revertPQ = make(PriorityQueue, 0)
    heap.Init(&sched.revertPQ)
    sched.entryPoint = entryPoint
    sched.JobLocker = new(sync.Mutex)
    sched.PQLocker = new(sync.Mutex)
    sched.FuncLocker = new(sync.Mutex)
    sched.stats = make(map[string]*FuncStat)
    sched.driver = driver
    sched.jobPQ = make(map[string]*PriorityQueue)
    sched.timeout = timeout
    sched.alive = true
    return sched
}

func (sched *Sched) Serve() {
    parts := strings.SplitN(sched.entryPoint, "://", 2)
    if parts[0] == "unix" {
        sockCheck(parts[1])
    }
    sched.loadJobQueue()
    go sched.handleJobPQ()
    go sched.handleRevertPQ()
    listen, err := net.Listen(parts[0], parts[1])
    if err != nil {
        log.Fatal(err)
    }
    defer listen.Close()
    log.Printf("Periodic task system started on %s\n", sched.entryPoint)
    for {
        if !sched.alive {
            break
        }
        conn, err := listen.Accept()
        if err != nil {
            log.Fatal(err)
        }
        if sched.timeout > 0 {
            conn.SetDeadline(time.Now().Add(sched.timeout * time.Second))
        }
        sched.HandleConnection(conn)
    }
}

func (sched *Sched) NotifyJobTimer() {
    sched.jobTimer.Reset(time.Millisecond)
}

func (sched *Sched) NotifyRevertTimer() {
    sched.revTimer.Reset(time.Millisecond)
}

func (sched *Sched) HandleConnection(conn net.Conn) {
    c := Conn{Conn: conn}
    payload, err := c.Receive()
    if err != nil {
        return
    }
    switch protocol.ClientType(payload[0]) {
    case protocol.TYPE_CLIENT:
        client := NewClient(sched, c)
        go client.Handle()
        break
    case protocol.TYPE_WORKER:
        worker := NewWorker(sched, c)
        go worker.Handle()
        break
    default:
        log.Printf("Unsupport client %d\n", payload[0])
        c.Close()
        break
    }
}

func (sched *Sched) Done(jobId int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    removeListJob(sched.procQueue, jobId)
    job, err := sched.driver.Get(jobId)
    if err == nil {
        sched.driver.Delete(jobId)
        sched.DecrStatJob(job)
        sched.DecrStatProc(job)
        sched.removeRevertPQ(job)
    }
    return
}

func (sched *Sched) isDoJob(job driver.Job) bool {
    for e := sched.procQueue.Front(); e != nil; e = e.Next() {
        chk := e.Value.(driver.Job)
        if chk.Id == job.Id {
            return true
        }
    }
    return false
}

func (sched *Sched) SubmitJob(grabItem GrabItem, job driver.Job) bool {
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    if job.Name == "" {
        sched.driver.Delete(job.Id)
        return true
    }
    if sched.isDoJob(job) {
        return true
    }
    if !grabItem.w.alive {
        return false
    }
    if err := grabItem.w.HandleDo(grabItem.msgId, job); err != nil {
        grabItem.w.alive = false
        return false
    }
    now := time.Now()
    current := int64(now.Unix())
    job.Status = driver.JOB_STATUS_PROC
    job.RunAt = current
    sched.driver.Save(&job)
    sched.IncrStatProc(job)
    sched.pushRevertPQ(job)
    sched.NotifyRevertTimer()
    sched.procQueue.PushBack(job)
    sched.grabQueue.Remove(grabItem)
    return true
}

func (sched *Sched) lessItem() (lessItem *Item) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    maybeItem := make(map[string]*Item)
    for Func, stat := range sched.stats {
        if stat.Worker == 0 {
            continue
        }
        pq, ok := sched.jobPQ[Func]
        if !ok || pq.Len() == 0 {
            continue
        }

        item := heap.Pop(pq).(*Item)

        maybeItem[Func] = item

    }

    if len(maybeItem) == 0 {
        return nil
    }

    var lessFunc string

    for Func, item := range maybeItem {
        if lessItem == nil {
            lessItem = item
            lessFunc = Func
            continue
        }
        if lessItem.priority > item.priority {
            lessItem = item
            lessFunc = Func
        }
    }

    for Func, item := range maybeItem {
        if Func == lessFunc {
            continue
        }
        pq := sched.jobPQ[Func]
        heap.Push(pq, item)
    }
    return
}

func (sched *Sched) handleJobPQ() {
    var current time.Time
    var timestamp int64
    for {
        if !sched.alive {
            break
        }
        if sched.grabQueue.Len() == 0 {
            sched.jobTimer.Reset(time.Minute)
            current =<-sched.jobTimer.C
            continue
        }

        lessItem := sched.lessItem()

        if lessItem == nil {
            sched.jobTimer.Reset(time.Minute)
            current =<-sched.jobTimer.C
            continue
        }

        schedJob, err := sched.driver.Get(lessItem.value)

        if err != nil {
            log.Printf("handleJobPQ error job: %d %v\n", lessItem.value, err)
            continue
        }

        timestamp = int64(time.Now().Unix())

        if schedJob.SchedAt > timestamp {
            sched.jobTimer.Reset(time.Second * time.Duration(schedJob.SchedAt - timestamp))
            current =<-sched.jobTimer.C
            timestamp = int64(current.Unix())
            if schedJob.SchedAt > timestamp {
                sched.pushJobPQ(schedJob)
                continue
            }
        }

        grabItem, err := sched.grabQueue.Get(schedJob.Func)
        if err == nil {
            if !sched.SubmitJob(grabItem, schedJob) {
                sched.pushJobPQ(schedJob)
            }
        } else {
            sched.pushJobPQ(schedJob)
        }
    }
}

func (sched *Sched) handleRevertPQ() {
    var current time.Time
    var timestamp int64
    for {
        if !sched.alive {
            break
        }
        if sched.revertPQ.Len() == 0 {
            sched.revTimer.Reset(time.Minute)
            current =<-sched.revTimer.C
            continue
        }

        sched.PQLocker.Lock()
        item := heap.Pop(&sched.revertPQ).(*Item)
        sched.PQLocker.Unlock()

        if item == nil {
            sched.revTimer.Reset(time.Minute)
            current =<-sched.revTimer.C
            continue
        }

        revertJob, err := sched.driver.Get(item.value)

        if err != nil {
            log.Printf("handleRevertPQ error: job: %d %v\n", item.value, err)
            continue
        }

        timestamp = int64(time.Now().Unix())

        if item.priority > timestamp {
            sched.revTimer.Reset(time.Second * time.Duration(item.priority - timestamp))
            current =<-sched.revTimer.C
            timestamp = int64(current.Unix())
            if item.priority > timestamp {
                sched.pushRevertPQ(revertJob)
                continue
            }
        }

        sched.DecrStatProc(revertJob)
        revertJob.Status = driver.JOB_STATUS_READY
        sched.driver.Save(&revertJob)
        sched.pushJobPQ(revertJob)
        removeListJob(sched.procQueue, revertJob.Id)
    }
}

func (sched *Sched) Fail(jobId int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    removeListJob(sched.procQueue, jobId)
    job, _ := sched.driver.Get(jobId)
    sched.DecrStatProc(job)
    sched.removeRevertPQ(job)
    job.Status = driver.JOB_STATUS_READY
    sched.driver.Save(&job)
    sched.pushJobPQ(job)
    return
}

func (sched *Sched) getFuncStat(Func string) *FuncStat {
    defer sched.FuncLocker.Unlock()
    sched.FuncLocker.Lock()
    stat, ok := sched.stats[Func]
    if !ok {
        stat = new(FuncStat)
        sched.stats[Func] = stat
    }
    return stat
}

func (sched *Sched) IncrStatFunc(Func string) {
    stat := sched.getFuncStat(Func)
    stat.Worker.Incr()
}

func (sched *Sched) DecrStatFunc(Func string) {
    stat := sched.getFuncStat(Func)
    stat.Worker.Decr()
}

func (sched *Sched) IncrStatJob(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    stat.Job.Incr()
}

func (sched *Sched) DecrStatJob(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    stat.Job.Decr()
}

func (sched *Sched) IncrStatProc(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    if job.Status == driver.JOB_STATUS_PROC {
        stat.Processing.Incr()
    }
}

func (sched *Sched) DecrStatProc(job driver.Job) {
    stat := sched.getFuncStat(job.Func)
    if job.Status == driver.JOB_STATUS_PROC {
        stat.Processing.Decr()
    }
}

func (sched *Sched) SchedLater(jobId int64, delay int64) {
    defer sched.NotifyJobTimer()
    defer sched.NotifyRevertTimer()
    defer sched.JobLocker.Unlock()
    sched.JobLocker.Lock()
    removeListJob(sched.procQueue, jobId)
    job, _ := sched.driver.Get(jobId)
    sched.DecrStatProc(job)
    sched.removeRevertPQ(job)
    job.Status = driver.JOB_STATUS_READY
    var now = time.Now()
    job.SchedAt = int64(now.Unix()) + delay
    sched.driver.Save(&job)
    sched.pushJobPQ(job)
    return
}

func (sched *Sched) pushJobPQ(job driver.Job) bool {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_READY {
        pq, ok := sched.jobPQ[job.Func]
        if !ok {
            pq1 := make(PriorityQueue, 0)
            pq = &pq1
            sched.jobPQ[job.Func] = pq
            heap.Init(pq)
        }
        item := &Item{
            value: job.Id,
            priority: job.SchedAt,
        }
        heap.Push(pq, item)
        return true
    }
    return false
}

func (sched *Sched) pushRevertPQ(job driver.Job) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_PROC && job.Timeout > 0 {
        runAt := job.RunAt
        if runAt == 0 {
            runAt = job.SchedAt
        }
        item := &Item{
            value: job.Id,
            priority: runAt + job.Timeout,
        }
        heap.Push(&sched.revertPQ, item)
    }
}

func (sched *Sched) removeRevertPQ(job driver.Job) {
    defer sched.PQLocker.Unlock()
    sched.PQLocker.Lock()
    if job.Status == driver.JOB_STATUS_PROC && job.Timeout > 0 {
        for _, item := range sched.revertPQ {
            if item.value == job.Id {
                heap.Remove(&sched.revertPQ, item.index)
                break
            }
        }
    }
}

func (sched *Sched) loadJobQueue() {
    updateQueue := make([]driver.Job, 0)
    removeQueue := make([]driver.Job, 0)
    var now = time.Now()
    current := int64(now.Unix())

    iter := sched.driver.NewIterator(nil)
    for {
        if !iter.Next() {
            break
        }
        job := iter.Value()
        if job.Name == "" {
            removeQueue = append(removeQueue, job)
            continue
        }
        sched.IncrStatJob(job)
        sched.pushJobPQ(job)
        runAt := job.RunAt
        if runAt < job.SchedAt {
            runAt = job.SchedAt
        }
        if runAt + job.Timeout < current {
            updateQueue = append(updateQueue, job)
        } else {
            sched.procQueue.PushBack(job)
            sched.IncrStatProc(job)
            sched.pushRevertPQ(job)
        }
    }

    iter.Close()

    for _, job := range updateQueue {
        job.Status = driver.JOB_STATUS_READY
        sched.driver.Save(&job)
    }

    for _, job := range removeQueue {
        sched.driver.Delete(job.Id)
    }
}

func (sched *Sched) Close() {
    sched.alive = false
    sched.driver.Close()
    log.Printf("Periodic task system shutdown\n")
}
<|endoftext|>


<commit_before>package sched

import (
	pq "github.com/cenkalti/gopqueue"
	"runtime"
	"sync"
	"time"
)

type Event struct {
	time		time.Time
	action		func ()
}

func (e *Event) Less(other interface {}) bool {
	return e.time.Before(other.(*Event).time)
}

type Scheduler struct {
	queue		*pq.Queue
	lock		sync.RWMutex
}

func New() *Scheduler {
	return &Scheduler{
		queue:    pq.New(0),
	}
}

func (s *Scheduler) EnterAbs(time time.Time, action func ()) Event {
	s.lock.Lock()
	defer s.lock.Unlock()

	event := Event{time, action}
	s.queue.Enqueue(&event)
	return event
}

func (s *Scheduler) Enter(delay time.Duration, action func ()) Event {
	diff := time.Now().Add(delay)
	return s.EnterAbs(diff, action)
}

func (s *Scheduler) Empty() bool {
	s.lock.RLock()
	defer s.lock.RUnlock()

	return s.queue.IsEmpty()
}

func (s *Scheduler) Len() int {
	s.lock.RLock()
	defer s.lock.RUnlock()

	return s.queue.Len()
}

func (s *Scheduler) Run() {
	var delay bool

	for {
		s.lock.Lock()
		min := s.queue.Peek()
		if min == nil {
			s.lock.Unlock()
			break
		}

		event := min.(*Event)
		now := time.Now()

		if event.time.After(now) {
			delay = true
		} else {
			delay = false
			s.queue.Dequeue()
		}
		s.lock.Unlock()

		if delay == true {
			time.Sleep(event.time.Sub(now))
		} else {
			go event.action()
			runtime.Gosched() // Don't know if this is required
		}
	}
}
<commit_msg>add method comments<commit_after>package sched

import (
	pq "github.com/cenkalti/gopqueue"
	"runtime"
	"sync"
	"time"
)

type Event struct {
	time		time.Time
	action		func ()
}

func (e *Event) Less(other interface {}) bool {
	return e.time.Before(other.(*Event).time)
}

type Scheduler struct {
	queue		*pq.Queue
	lock		sync.RWMutex
}

// New creates a new scheduler and return it's pointer.
func New() *Scheduler {
	return &Scheduler{
		queue:    pq.New(0),
	}
}

// EnterAbs adds a new event to the queue at an absolute time.
func (s *Scheduler) EnterAbs(time time.Time, action func ()) Event {
	s.lock.Lock()
	defer s.lock.Unlock()

	event := Event{time, action}
	s.queue.Enqueue(&event)
	return event
}

// Enter adds an new event to the queue to run after delay.
func (s *Scheduler) Enter(delay time.Duration, action func ()) Event {
	diff := time.Now().Add(delay)
	return s.EnterAbs(diff, action)
}

// Empty returns true if there is not events in the queue.
func (s *Scheduler) Empty() bool {
	s.lock.RLock()
	defer s.lock.RUnlock()

	return s.queue.IsEmpty()
}

// Len returns the number of items in the scheduler's event queue.
func (s *Scheduler) Len() int {
	s.lock.RLock()
	defer s.lock.RUnlock()

	return s.queue.Len()
}

// Run executes events until the queue is empty.
func (s *Scheduler) Run() {
	var delay bool

	for {
		s.lock.Lock()
		min := s.queue.Peek()
		if min == nil {
			s.lock.Unlock()
			break
		}

		event := min.(*Event)
		now := time.Now()

		if event.time.After(now) {
			delay = true
		} else {
			delay = false
			s.queue.Dequeue()
		}
		s.lock.Unlock()

		if delay == true {
			time.Sleep(event.time.Sub(now))
		} else {
			go event.action()
			runtime.Gosched() // Don't know if this is required
		}
	}
}
<|endoftext|>


<commit_before>package gorm

import (
	"errors"
	"fmt"
	"go/ast"
	"strings"
	"time"

	"reflect"
	"regexp"
)

type Scope struct {
	Value         interface{}
	indirectValue *reflect.Value
	Search        *search
	Sql           string
	SqlVars       []interface{}
	db            *DB
	_values       map[string]interface{}
	skipLeft      bool
	primaryKey    string
}

func (scope *Scope) IndirectValue() reflect.Value {
	if scope.indirectValue == nil {
		value := reflect.Indirect(reflect.ValueOf(scope.Value))
		scope.indirectValue = &value
	}
	return *scope.indirectValue
}

// NewScope create scope for callbacks, including DB's search information
func (db *DB) NewScope(value interface{}) *Scope {
	db.Value = value
	return &Scope{db: db, Search: db.search, Value: value, _values: map[string]interface{}{}}
}

// New create a new Scope without search information
func (scope *Scope) New(value interface{}) *Scope {
	return &Scope{db: scope.db.parent, Search: &search{}, Value: value}
}

// NewDB create a new DB without search information
func (scope *Scope) NewDB() *DB {
	return scope.db.new()
}

// DB get *sql.DB
func (scope *Scope) DB() sqlCommon {
	return scope.db.db
}

// SkipLeft skip remaining callbacks
func (scope *Scope) SkipLeft() {
	scope.skipLeft = true
}

// Quote used to quote database column name according to database dialect
func (scope *Scope) Quote(str string) string {
	return scope.Dialect().Quote(str)
}

// Dialect get dialect
func (scope *Scope) Dialect() Dialect {
	return scope.db.parent.dialect
}

// Err write error
func (scope *Scope) Err(err error) error {
	if err != nil {
		scope.db.err(err)
	}
	return err
}

// Log print log message
func (scope *Scope) Log(v ...interface{}) {
	scope.db.log(v...)
}

// HasError check if there are any error
func (scope *Scope) HasError() bool {
	return scope.db.Error != nil
}

// PrimaryKey get the primary key's column name
func (scope *Scope) PrimaryKey() string {
	if scope.primaryKey != "" {
		return scope.primaryKey
	}

	scope.primaryKey = ToSnake(GetPrimaryKey(scope.Value))
	return scope.primaryKey
}

// PrimaryKeyZero check the primary key is blank or not
func (scope *Scope) PrimaryKeyZero() bool {
	return isBlank(reflect.ValueOf(scope.PrimaryKeyValue()))
}

// PrimaryKeyValue get the primary key's value
func (scope *Scope) PrimaryKeyValue() interface{} {
	if scope.IndirectValue().Kind() == reflect.Struct {
		if field := scope.IndirectValue().FieldByName(SnakeToUpperCamel(scope.PrimaryKey())); field.IsValid() {
			return field.Interface()
		}
	}
	return 0
}

// HasColumn to check if has column
func (scope *Scope) HasColumn(name string) bool {
	_, result := scope.FieldByName(name)
	return result
}

// FieldByName to get column's value and existence
func (scope *Scope) FieldByName(name string) (interface{}, bool) {
	return FieldByName(name, scope.Value)
}

// SetColumn to set the column's value
func (scope *Scope) SetColumn(column string, value interface{}) {
	if scope.Value == nil {
		return
	}

	setFieldValue(scope.IndirectValue().FieldByName(SnakeToUpperCamel(column)), value)
}

// CallMethod invoke method with necessary argument
func (scope *Scope) CallMethod(name string) {
	if scope.Value == nil {
		return
	}

	call := func(value interface{}) {
		if fm := reflect.ValueOf(value).MethodByName(name); fm.IsValid() {
			fi := fm.Interface()
			if f, ok := fi.(func()); ok {
				f()
			} else if f, ok := fi.(func(s *Scope)); ok {
				f(scope)
			} else if f, ok := fi.(func(s *DB)); ok {
				f(scope.db.new())
			} else if f, ok := fi.(func() error); ok {
				scope.Err(f())
			} else if f, ok := fi.(func(s *Scope) error); ok {
				scope.Err(f(scope))
			} else if f, ok := fi.(func(s *DB) error); ok {
				scope.Err(f(scope.db.new()))
			} else {
				scope.Err(errors.New(fmt.Sprintf("unsupported function %v", name)))
			}
		}
	}

	if values := scope.IndirectValue(); values.Kind() == reflect.Slice {
		for i := 0; i < values.Len(); i++ {
			call(values.Index(i).Addr().Interface())
		}
	} else {
		call(scope.Value)
	}
}

// AddToVars add value as sql's vars, gorm will escape them
func (scope *Scope) AddToVars(value interface{}) string {
	scope.SqlVars = append(scope.SqlVars, value)
	return scope.Dialect().BinVar(len(scope.SqlVars))
}

// TableName get table name
var pluralMapKeys = []*regexp.Regexp{regexp.MustCompile("ch$"), regexp.MustCompile("ss$"), regexp.MustCompile("sh$"), regexp.MustCompile("day$"), regexp.MustCompile("y$"), regexp.MustCompile("x$"), regexp.MustCompile("([^s])s?$")}
var pluralMapValues = []string{"ches", "sses", "shes", "days", "ies", "xes", "${1}s"}

func (scope *Scope) TableName() string {
	if scope.Search != nil && len(scope.Search.TableName) > 0 {
		return scope.Search.TableName
	} else {
		if scope.Value == nil {
			scope.Err(errors.New("can't get table name"))
			return ""
		}

		data := scope.IndirectValue()
		if data.Kind() == reflect.Slice {
			elem := data.Type().Elem()
			if elem.Kind() == reflect.Ptr {
				elem = elem.Elem()
			}
			data = reflect.New(elem).Elem()
		}

		if fm := data.MethodByName("TableName"); fm.IsValid() {
			if v := fm.Call([]reflect.Value{}); len(v) > 0 {
				if result, ok := v[0].Interface().(string); ok {
					return result
				}
			}
		}

		str := ToSnake(data.Type().Name())

		if !scope.db.parent.singularTable {
			for index, reg := range pluralMapKeys {
				if reg.MatchString(str) {
					return reg.ReplaceAllString(str, pluralMapValues[index])
				}
			}
		}

		return str
	}
}

func (scope *Scope) QuotedTableName() string {
	if scope.Search != nil && len(scope.Search.TableName) > 0 {
		return scope.Search.TableName
	} else {
		keys := strings.Split(scope.TableName(), ".")
		for i, v := range keys {
			keys[i] = scope.Quote(v)
		}
		return strings.Join(keys, ".")
	}
}

// CombinedConditionSql get combined condition sql
func (scope *Scope) CombinedConditionSql() string {
	return scope.joinsSql() + scope.whereSql() + scope.groupSql() +
		scope.havingSql() + scope.orderSql() + scope.limitSql() + scope.offsetSql()
}

func (scope *Scope) fieldFromStruct(fieldStruct reflect.StructField) *Field {
	var field Field
	field.Name = fieldStruct.Name
	field.DBName = ToSnake(fieldStruct.Name)

	value := scope.IndirectValue().FieldByName(fieldStruct.Name)
	indirectValue := reflect.Indirect(value)
	field.Value = value.Interface()
	field.IsBlank = isBlank(value)

	// Search for primary key tag identifier
	settings := parseTagSetting(fieldStruct.Tag.Get("gorm"))
	if _, ok := settings["PRIMARY_KEY"]; scope.PrimaryKey() == field.DBName || ok {
		field.isPrimaryKey = true
	}

	if field.isPrimaryKey {
		scope.primaryKey = field.DBName
	}

	if scope.db != nil {
		field.Tag = fieldStruct.Tag
		field.SqlTag = scope.sqlTagForField(&field)

		// parse association
		typ := indirectValue.Type()
		foreignKey := SnakeToUpperCamel(settings["FOREIGNKEY"])
		associationForeignKey := SnakeToUpperCamel(settings["ASSOCIATIONFOREIGNKEY"])
		many2many := settings["MANY2MANY"]
		scopeTyp := scope.IndirectValue().Type()

		switch indirectValue.Kind() {
		case reflect.Slice:
			typ = typ.Elem()

			if typ.Kind() == reflect.Struct {
				if foreignKey == "" {
					foreignKey = scopeTyp.Name() + "Id"
				}
				if associationForeignKey == "" {
					associationForeignKey = typ.Name() + "Id"
				}

				// if not many to many, foreign key could be null
				if many2many == "" {
					if !reflect.New(typ).Elem().FieldByName(foreignKey).IsValid() {
						foreignKey = ""
					}
				}

				field.AfterAssociation = true
				field.JoinTable = &joinTable{
					joinTable:             many2many,
					foreignKey:            foreignKey,
					associationForeignKey: associationForeignKey,
				}
			}
		case reflect.Struct:
			if !field.IsTime() && !field.IsScanner() {
				if foreignKey == "" && scope.HasColumn(field.Name+"Id") {
					field.JoinTable = &joinTable{foreignKey: field.Name + "Id"}
					field.BeforeAssociation = true
				} else if scope.HasColumn(foreignKey) {
					field.JoinTable = &joinTable{foreignKey: foreignKey}
					field.BeforeAssociation = true
				} else {
					if foreignKey == "" {
						foreignKey = scopeTyp.Name() + "Id"
					}
					if reflect.New(typ).Elem().FieldByName(foreignKey).IsValid() {
						field.JoinTable = &joinTable{foreignKey: foreignKey}
					}
					field.AfterAssociation = true
				}
			}
		}
	}
	return &field
}

// Fields get value's fields
func (scope *Scope) Fields() []*Field {
	indirectValue := scope.IndirectValue()
	fields := []*Field{}

	if !indirectValue.IsValid() {
		return fields
	}

	scopeTyp := indirectValue.Type()
	for i := 0; i < scopeTyp.NumField(); i++ {
		fieldStruct := scopeTyp.Field(i)
		if !ast.IsExported(fieldStruct.Name) {
			continue
		}
		fields = append(fields, scope.fieldFromStruct(fieldStruct))
	}

	return fields
}

// Raw set sql
func (scope *Scope) Raw(sql string) *Scope {
	scope.Sql = strings.Replace(sql, "$$", "?", -1)
	return scope
}

// Exec invoke sql
func (scope *Scope) Exec() *Scope {
	defer scope.Trace(time.Now())

	if !scope.HasError() {
		result, err := scope.DB().Exec(scope.Sql, scope.SqlVars...)
		if scope.Err(err) == nil {
			if count, err := result.RowsAffected(); err == nil {
				scope.db.RowsAffected = count
			}
		}
	}
	return scope
}

// Set set value by name
func (scope *Scope) Set(name string, value interface{}) *Scope {
	scope._values[name] = value
	return scope
}

// Get get value by name
func (scope *Scope) Get(name string) (value interface{}, ok bool) {
	value, ok = scope._values[name]
	return
}

// Trace print sql log
func (scope *Scope) Trace(t time.Time) {
	if len(scope.Sql) > 0 {
		scope.db.slog(scope.Sql, t, scope.SqlVars...)
	}
}

// Begin start a transaction
func (scope *Scope) Begin() *Scope {
	if db, ok := scope.DB().(sqlDb); ok {
		if tx, err := db.Begin(); err == nil {
			scope.db.db = interface{}(tx).(sqlCommon)
			scope.Set("gorm:started_transaction", true)
		}
	}
	return scope
}

// CommitOrRollback commit current transaction if there is no error, otherwise rollback it
func (scope *Scope) CommitOrRollback() *Scope {
	if _, ok := scope.Get("gorm:started_transaction"); ok {
		if db, ok := scope.db.db.(sqlTx); ok {
			if scope.HasError() {
				db.Rollback()
			} else {
				db.Commit()
			}
			scope.db.db = scope.db.parent.db
		}
	}
	return scope
}
<commit_msg>Refact scope Fields<commit_after>package gorm

import (
	"errors"
	"fmt"
	"go/ast"
	"strings"
	"time"

	"reflect"
	"regexp"
)

type Scope struct {
	Value         interface{}
	indirectValue *reflect.Value
	Search        *search
	Sql           string
	SqlVars       []interface{}
	db            *DB
	_values       map[string]interface{}
	skipLeft      bool
	primaryKey    string
}

func (scope *Scope) IndirectValue() reflect.Value {
	if scope.indirectValue == nil {
		value := reflect.Indirect(reflect.ValueOf(scope.Value))
		scope.indirectValue = &value
	}
	return *scope.indirectValue
}

// NewScope create scope for callbacks, including DB's search information
func (db *DB) NewScope(value interface{}) *Scope {
	db.Value = value
	return &Scope{db: db, Search: db.search, Value: value, _values: map[string]interface{}{}}
}

// New create a new Scope without search information
func (scope *Scope) New(value interface{}) *Scope {
	return &Scope{db: scope.db.parent, Search: &search{}, Value: value}
}

// NewDB create a new DB without search information
func (scope *Scope) NewDB() *DB {
	return scope.db.new()
}

// DB get *sql.DB
func (scope *Scope) DB() sqlCommon {
	return scope.db.db
}

// SkipLeft skip remaining callbacks
func (scope *Scope) SkipLeft() {
	scope.skipLeft = true
}

// Quote used to quote database column name according to database dialect
func (scope *Scope) Quote(str string) string {
	return scope.Dialect().Quote(str)
}

// Dialect get dialect
func (scope *Scope) Dialect() Dialect {
	return scope.db.parent.dialect
}

// Err write error
func (scope *Scope) Err(err error) error {
	if err != nil {
		scope.db.err(err)
	}
	return err
}

// Log print log message
func (scope *Scope) Log(v ...interface{}) {
	scope.db.log(v...)
}

// HasError check if there are any error
func (scope *Scope) HasError() bool {
	return scope.db.Error != nil
}

// PrimaryKey get the primary key's column name
func (scope *Scope) PrimaryKey() string {
	if scope.primaryKey != "" {
		return scope.primaryKey
	}

	scope.primaryKey = ToSnake(GetPrimaryKey(scope.Value))
	return scope.primaryKey
}

// PrimaryKeyZero check the primary key is blank or not
func (scope *Scope) PrimaryKeyZero() bool {
	return isBlank(reflect.ValueOf(scope.PrimaryKeyValue()))
}

// PrimaryKeyValue get the primary key's value
func (scope *Scope) PrimaryKeyValue() interface{} {
	if scope.IndirectValue().Kind() == reflect.Struct {
		if field := scope.IndirectValue().FieldByName(SnakeToUpperCamel(scope.PrimaryKey())); field.IsValid() {
			return field.Interface()
		}
	}
	return 0
}

// HasColumn to check if has column
func (scope *Scope) HasColumn(name string) bool {
	_, result := scope.FieldByName(name)
	return result
}

// FieldByName to get column's value and existence
func (scope *Scope) FieldByName(name string) (interface{}, bool) {
	return FieldByName(name, scope.Value)
}

// SetColumn to set the column's value
func (scope *Scope) SetColumn(column string, value interface{}) {
	if scope.Value == nil {
		return
	}

	setFieldValue(scope.IndirectValue().FieldByName(SnakeToUpperCamel(column)), value)
}

// CallMethod invoke method with necessary argument
func (scope *Scope) CallMethod(name string) {
	if scope.Value == nil {
		return
	}

	call := func(value interface{}) {
		if fm := reflect.ValueOf(value).MethodByName(name); fm.IsValid() {
			fi := fm.Interface()
			if f, ok := fi.(func()); ok {
				f()
			} else if f, ok := fi.(func(s *Scope)); ok {
				f(scope)
			} else if f, ok := fi.(func(s *DB)); ok {
				f(scope.db.new())
			} else if f, ok := fi.(func() error); ok {
				scope.Err(f())
			} else if f, ok := fi.(func(s *Scope) error); ok {
				scope.Err(f(scope))
			} else if f, ok := fi.(func(s *DB) error); ok {
				scope.Err(f(scope.db.new()))
			} else {
				scope.Err(errors.New(fmt.Sprintf("unsupported function %v", name)))
			}
		}
	}

	if values := scope.IndirectValue(); values.Kind() == reflect.Slice {
		for i := 0; i < values.Len(); i++ {
			call(values.Index(i).Addr().Interface())
		}
	} else {
		call(scope.Value)
	}
}

// AddToVars add value as sql's vars, gorm will escape them
func (scope *Scope) AddToVars(value interface{}) string {
	scope.SqlVars = append(scope.SqlVars, value)
	return scope.Dialect().BinVar(len(scope.SqlVars))
}

// TableName get table name
var pluralMapKeys = []*regexp.Regexp{regexp.MustCompile("ch$"), regexp.MustCompile("ss$"), regexp.MustCompile("sh$"), regexp.MustCompile("day$"), regexp.MustCompile("y$"), regexp.MustCompile("x$"), regexp.MustCompile("([^s])s?$")}
var pluralMapValues = []string{"ches", "sses", "shes", "days", "ies", "xes", "${1}s"}

func (scope *Scope) TableName() string {
	if scope.Search != nil && len(scope.Search.TableName) > 0 {
		return scope.Search.TableName
	} else {
		if scope.Value == nil {
			scope.Err(errors.New("can't get table name"))
			return ""
		}

		data := scope.IndirectValue()
		if data.Kind() == reflect.Slice {
			elem := data.Type().Elem()
			if elem.Kind() == reflect.Ptr {
				elem = elem.Elem()
			}
			data = reflect.New(elem).Elem()
		}

		if fm := data.MethodByName("TableName"); fm.IsValid() {
			if v := fm.Call([]reflect.Value{}); len(v) > 0 {
				if result, ok := v[0].Interface().(string); ok {
					return result
				}
			}
		}

		str := ToSnake(data.Type().Name())

		if !scope.db.parent.singularTable {
			for index, reg := range pluralMapKeys {
				if reg.MatchString(str) {
					return reg.ReplaceAllString(str, pluralMapValues[index])
				}
			}
		}

		return str
	}
}

func (scope *Scope) QuotedTableName() string {
	if scope.Search != nil && len(scope.Search.TableName) > 0 {
		return scope.Search.TableName
	} else {
		keys := strings.Split(scope.TableName(), ".")
		for i, v := range keys {
			keys[i] = scope.Quote(v)
		}
		return strings.Join(keys, ".")
	}
}

// CombinedConditionSql get combined condition sql
func (scope *Scope) CombinedConditionSql() string {
	return scope.joinsSql() + scope.whereSql() + scope.groupSql() +
		scope.havingSql() + scope.orderSql() + scope.limitSql() + scope.offsetSql()
}

func (scope *Scope) fieldFromStruct(fieldStruct reflect.StructField) *Field {
	var field Field
	field.Name = fieldStruct.Name
	field.DBName = ToSnake(fieldStruct.Name)

	value := scope.IndirectValue().FieldByName(fieldStruct.Name)
	indirectValue := reflect.Indirect(value)
	field.Value = value.Interface()
	field.IsBlank = isBlank(value)

	// Search for primary key tag identifier
	settings := parseTagSetting(fieldStruct.Tag.Get("gorm"))
	if _, ok := settings["PRIMARY_KEY"]; scope.PrimaryKey() == field.DBName || ok {
		field.isPrimaryKey = true
	}

	if field.isPrimaryKey {
		scope.primaryKey = field.DBName
	}

	if scope.db != nil {
		field.Tag = fieldStruct.Tag
		field.SqlTag = scope.sqlTagForField(&field)

		// parse association
		typ := indirectValue.Type()
		foreignKey := SnakeToUpperCamel(settings["FOREIGNKEY"])
		associationForeignKey := SnakeToUpperCamel(settings["ASSOCIATIONFOREIGNKEY"])
		many2many := settings["MANY2MANY"]
		scopeTyp := scope.IndirectValue().Type()

		switch indirectValue.Kind() {
		case reflect.Slice:
			typ = typ.Elem()

			if typ.Kind() == reflect.Struct {
				if foreignKey == "" {
					foreignKey = scopeTyp.Name() + "Id"
				}
				if associationForeignKey == "" {
					associationForeignKey = typ.Name() + "Id"
				}

				// if not many to many, foreign key could be null
				if many2many == "" {
					if !reflect.New(typ).Elem().FieldByName(foreignKey).IsValid() {
						foreignKey = ""
					}
				}

				field.AfterAssociation = true
				field.JoinTable = &joinTable{
					joinTable:             many2many,
					foreignKey:            foreignKey,
					associationForeignKey: associationForeignKey,
				}
			}
		case reflect.Struct:
			if !field.IsTime() && !field.IsScanner() {
				if foreignKey == "" && scope.HasColumn(field.Name+"Id") {
					field.JoinTable = &joinTable{foreignKey: field.Name + "Id"}
					field.BeforeAssociation = true
				} else if scope.HasColumn(foreignKey) {
					field.JoinTable = &joinTable{foreignKey: foreignKey}
					field.BeforeAssociation = true
				} else {
					if foreignKey == "" {
						foreignKey = scopeTyp.Name() + "Id"
					}
					if reflect.New(typ).Elem().FieldByName(foreignKey).IsValid() {
						field.JoinTable = &joinTable{foreignKey: foreignKey}
					}
					field.AfterAssociation = true
				}
			}
		}
	}
	return &field
}

// Fields get value's fields
func (scope *Scope) Fields() (fields []*Field) {
	if scope.IndirectValue().IsValid() {
		scopeTyp := scope.IndirectValue().Type()
		for i := 0; i < scopeTyp.NumField(); i++ {
			fieldStruct := scopeTyp.Field(i)
			if !ast.IsExported(fieldStruct.Name) {
				continue
			}
			fields = append(fields, scope.fieldFromStruct(fieldStruct))
		}
	}
	return
}

// Raw set sql
func (scope *Scope) Raw(sql string) *Scope {
	scope.Sql = strings.Replace(sql, "$$", "?", -1)
	return scope
}

// Exec invoke sql
func (scope *Scope) Exec() *Scope {
	defer scope.Trace(time.Now())

	if !scope.HasError() {
		result, err := scope.DB().Exec(scope.Sql, scope.SqlVars...)
		if scope.Err(err) == nil {
			if count, err := result.RowsAffected(); err == nil {
				scope.db.RowsAffected = count
			}
		}
	}
	return scope
}

// Set set value by name
func (scope *Scope) Set(name string, value interface{}) *Scope {
	scope._values[name] = value
	return scope
}

// Get get value by name
func (scope *Scope) Get(name string) (value interface{}, ok bool) {
	value, ok = scope._values[name]
	return
}

// Trace print sql log
func (scope *Scope) Trace(t time.Time) {
	if len(scope.Sql) > 0 {
		scope.db.slog(scope.Sql, t, scope.SqlVars...)
	}
}

// Begin start a transaction
func (scope *Scope) Begin() *Scope {
	if db, ok := scope.DB().(sqlDb); ok {
		if tx, err := db.Begin(); err == nil {
			scope.db.db = interface{}(tx).(sqlCommon)
			scope.Set("gorm:started_transaction", true)
		}
	}
	return scope
}

// CommitOrRollback commit current transaction if there is no error, otherwise rollback it
func (scope *Scope) CommitOrRollback() *Scope {
	if _, ok := scope.Get("gorm:started_transaction"); ok {
		if db, ok := scope.db.db.(sqlTx); ok {
			if scope.HasError() {
				db.Rollback()
			} else {
				db.Commit()
			}
			scope.db.db = scope.db.parent.db
		}
	}
	return scope
}
<|endoftext|>


<commit_before>package main

import (
    "fmt"
    "runtime"
)

func main() {
    //Run with max possible CPUs
    parallelism := 4
    runtime.GOMAXPROCS(runtime.NumCPU())

    var N int = 1e3 //Number of integers that will be in the channel

    c := make(chan int, parallelism) //Create our channel

    //In a separate thread, fill the channel
    go func() {
        defer close(c)
        fillChannel(c, N)
    }()

    //In this thread, drain the channel
    sum := drainChannel(c)

    fmt.Println("By manual counting, the sum of all integers from 0 to",(N-1),"is",sum)
}

func fillChannel(c chan int, N int) {
    for i := 0; i < N; i++ {
        c <- i
    }
}

func drainChannel(c chan int) int {
    x := 0
    for i := range c {
        x += i
    }

    return x
}<commit_msg>Increasing channel buffer size for faster execution.<commit_after>package main

import (
    "fmt"
    "runtime"
)

func main() {
    //Run with max possible CPUs
    numCpu := runtime.NumCPU()
    runtime.GOMAXPROCS(numCpu)

    //Number of integers that will be in the channel
    var N int = 1e7

    //Create our channel with enough space to hold all integers at once
    c := make(chan int, N)

    //In a separate thread, fill the channel
    go func() {
        defer close(c)
        start := 0
        stop := N
        fillChannel(c, start, stop)
    }()

    //In this thread, drain the channel
    sum := drainChannel(c)

    fmt.Println("By manual counting, the sum of all integers from 0 to",(N-1),"is",sum)
}

func fillChannel(c chan int, start int, stop int) {
    for i := start; i < stop; i++ {
        c <- i
    }
}

func drainChannel(c chan int) int64 {
    var x int64 = 0
    for i := range c {
        x += int64(i)
    }

    return x
}<|endoftext|>


<commit_before>package txtdirect

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"strings"
)

// fallback redirects the request to the given fallback address
// and if it's not provided it will check txtdirect config for
// default fallback address
func fallback(w http.ResponseWriter, r *http.Request, fallback, recordType, fallbackType string, code int, c Config) {
	if code == http.StatusMovedPermanently {
		w.Header().Add("Cache-Control", fmt.Sprintf("max-age=%d", status301CacheAge))
	}
	w.Header().Add("Status-Code", strconv.Itoa(code))

	if fallbackType != "global" {

		// Fetch records from request's context and set the []record type on them
		records := r.Context().Value("records").([]record)

		// Redirect to first record's `to=` field
		if fallbackType == "to" && records[0].To != "" {
			http.Redirect(w, r, records[0].To, code)
			if c.Prometheus.Enable {
				countFallback(r, records[0].Type, fallbackType, code)
			}
		}

		// Redirect to first record's `website=` field
		if fallbackType == "website" && records[0].Website != "" {
			http.Redirect(w, r, records[0].Website, code)
			if c.Prometheus.Enable {
				countFallback(r, records[0].Type, fallbackType, code)
			}
		}

		// Dockerv2 root fallback
		if fallbackType == "root" && records[0].Root != "" {
			http.Redirect(w, r, records[0].Root, code)
			if c.Prometheus.Enable {
				countFallback(r, records[len(records)-1].Type, fallbackType, code)
			}
		}

		// Special case when path is used in fetching the final record
		var pathRecord record
		if len(records) >= 2 {
			pathRecord = records[len(records)-2]
		}

		if fallbackType == "root" && pathRecord.Root != "" {
			http.Redirect(w, r, pathRecord.Root, code)
			if c.Prometheus.Enable {
				countFallback(r, records[len(records)-1].Type, fallbackType, code)
			}
		}

		if pathRecord.To != "" {
			http.Redirect(w, r, pathRecord.To, code)
			if c.Prometheus.Enable {
				countFallback(r, records[len(records)-1].Type, fallbackType, code)
			}
		}

		// If non of the above cases applied on the record, jump into global redirects
		globalFallbacks(w, r, c, code, records[len(records)-1].Type)
		log.Printf("[txtdirect]: %s > %s", r.Host+r.URL.Path, w.Header().Get("Location"))
		return
	}

	globalFallbacks(w, r, c, code, "")

	log.Printf("[txtdirect]: %s > %s", r.Host+r.URL.Path, w.Header().Get("Location"))
}

func addRecordToContext(r *http.Request, rec record) *http.Request {
	// Fetch fallback config from context and add the record to it
	recordsContext := r.Context().Value("records")

	// Create a new records field in the context if it doesn't exist
	if recordsContext == nil {
		return r.WithContext(context.WithValue(r.Context(), "records", []record{rec}))
	}

	records := append(recordsContext.([]record), rec)

	// Replace the fallback config instance inside the request's context
	return r.WithContext(context.WithValue(r.Context(), "records", records))
}

func countFallback(r *http.Request, recType, fallbackType string, code int) {
	FallbacksCount.WithLabelValues(r.Host, recType, fallbackType).Add(1)
	RequestsByStatus.WithLabelValues(r.URL.Host, string(code)).Add(1)
}

func globalFallbacks(w http.ResponseWriter, r *http.Request, c Config, code int, recordType string) {
	if contains(c.Enable, "www") {
		s := strings.Join([]string{defaultProtocol, "://", defaultSub, ".", r.URL.Host}, "")
		http.Redirect(w, r, s, code)
		if c.Prometheus.Enable {
			countFallback(r, recordType, "subdomain", code)
		}
	} else if c.Redirect != "" {
		w.Header().Set("Status-Code", strconv.Itoa(http.StatusMovedPermanently))

		http.Redirect(w, r, c.Redirect, http.StatusMovedPermanently)

		if c.Prometheus.Enable {
			countFallback(r, recordType, "redirect", http.StatusMovedPermanently)
		}
	} else {
		http.NotFound(w, r)
	}
}
<commit_msg>(fallback): Add fallback struct<commit_after>package txtdirect

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"strconv"
	"strings"
)

type Fallback struct {
	rw      http.ResponseWriter
	request *http.Request
	config  Config

	records    []record
	pathRecord record

	fallbackType string
	code         int
}

// fallback redirects the request to the given fallback address
// and if it's not provided it will check txtdirect config for
// default fallback address
func fallback(w http.ResponseWriter, r *http.Request, fallbackType string, code int, c Config) {
	if code == http.StatusMovedPermanently {
		w.Header().Add("Cache-Control", fmt.Sprintf("max-age=%d", status301CacheAge))
	}
	w.Header().Add("Status-Code", strconv.Itoa(code))

	f := Fallback{
		rw:           w,
		request:      r,
		config:       c,
		fallbackType: fallbackType,
		code:         code,
	}

	if fallbackType != "global" {

		// Fetch records from request's context and set the []record type on them
		f.fetchRecords()

		// Redirect to first record's `to=` field
		if fallbackType == "to" && f.records[0].To != "" {
			http.Redirect(w, r, f.records[0].To, code)
			f.countFallback(f.records[0].Type)
		}

		// Redirect to first record's `website=` field
		if fallbackType == "website" && f.records[0].Website != "" {
			http.Redirect(w, r, f.records[0].Website, code)
			f.countFallback(f.records[0].Type)
		}

		// Dockerv2 root fallback
		if fallbackType == "root" && f.records[0].Root != "" {
			http.Redirect(w, r, f.records[0].Root, code)
			f.countFallback(f.records[0].Type)
		}

		if fallbackType == "root" && f.pathRecord.Root != "" {
			http.Redirect(w, r, f.pathRecord.Root, code)
			f.countFallback(f.pathRecord.Type)
		}

		if f.pathRecord.To != "" {
			http.Redirect(w, r, f.pathRecord.To, code)
			f.countFallback(f.pathRecord.Type)
		}

		// If non of the above cases applied on the record, jump into global redirects
		f.globalFallbacks(f.records[len(f.records)-1].Type)
		log.Printf("[txtdirect]: %s > %s", r.Host+r.URL.Path, w.Header().Get("Location"))
		return
	}

	f.globalFallbacks("")

	log.Printf("[txtdirect]: %s > %s", r.Host+r.URL.Path, w.Header().Get("Location"))
}

func addRecordToContext(r *http.Request, rec record) *http.Request {
	// Fetch fallback config from context and add the record to it
	recordsContext := r.Context().Value("records")

	// Create a new records field in the context if it doesn't exist
	if recordsContext == nil {
		return r.WithContext(context.WithValue(r.Context(), "records", []record{rec}))
	}

	records := append(recordsContext.([]record), rec)

	// Replace the fallback config instance inside the request's context
	return r.WithContext(context.WithValue(r.Context(), "records", records))
}

func (f *Fallback) countFallback(recType string) {
	if f.config.Prometheus.Enable {
		FallbacksCount.WithLabelValues(f.request.Host, recType, f.fallbackType).Add(1)
		RequestsByStatus.WithLabelValues(f.request.URL.Host, string(f.code)).Add(1)
	}
}

func (f *Fallback) globalFallbacks(recordType string) {
	if contains(f.config.Enable, "www") {
		s := strings.Join([]string{defaultProtocol, "://", defaultSub, ".", f.request.URL.Host}, "")

		http.Redirect(f.rw, f.request, s, f.code)

		f.countFallback(recordType)
	} else if f.config.Redirect != "" {
		f.rw.Header().Set("Status-Code", strconv.Itoa(http.StatusMovedPermanently))

		http.Redirect(f.rw, f.request, f.config.Redirect, http.StatusMovedPermanently)

		f.code = http.StatusMovedPermanently

		f.countFallback(recordType)
	} else {
		http.NotFound(f.rw, f.request)
	}
}

func (f *Fallback) fetchRecords() {
	f.records = f.request.Context().Value("records").([]record)
	if len(f.records) >= 2 {
		f.pathRecord = f.records[len(f.records)-2]
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"database/sql"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path"
	"sync"

	_ "github.com/mattn/go-sqlite3"
)

const (
	thumbsDir      = "public/thumbs"
	bigThumbSize   = "1000x1000"
	smallThumbSize = "200x200"
	workers        = 4 // min: 1
)

func generateSmallThumb(photoPath, identifier string) (thumbPath string, err error) {
	thumbPath = path.Join(thumbsDir, fmt.Sprintf("%s_small.jpg", identifier))
	if _, err = os.Stat(thumbPath); os.IsNotExist(err) { // file does not exist
		err = exec.Command(
			"convert", photoPath,
			"-auto-orient",
			"-strip",
			"-resize", smallThumbSize+"^",
			"-gravity", "center",
			"-extent", smallThumbSize,
			thumbPath).Run()
	}
	return
}

func generateBigThumb(photoPath, identifier string) (thumbPath string, err error) {
	thumbPath = path.Join(thumbsDir, fmt.Sprintf("%s_big.jpg", identifier))
	if _, err = os.Stat(thumbPath); os.IsNotExist(err) { // file does not exist
		err = exec.Command(
			"convert", photoPath,
			"-auto-orient",
			"-strip",
			"-resize", bigThumbSize,
			thumbPath).Run()
	}
	return
}

func generateThumbsImpl(photoPath string) {
	identifier := fmt.Sprintf("%x", md5.Sum([]byte(photoPath)))

	bigThumbPath, err := generateBigThumb(photoPath, identifier)
	if err != nil { // error
		return
	}

	fmt.Println(bigThumbPath)

	smallThumbPath, err := generateSmallThumb(bigThumbPath, identifier)
	if err == nil { // success
		fmt.Println(smallThumbPath)
	}
}

func generateThumbs(ch chan string, wg *sync.WaitGroup) {
	defer wg.Done()

	for photoPath := range ch {
		generateThumbsImpl(photoPath)
	}
}

func main() {
	if workers < 1 {
		log.Fatal("number of workers must be at least 1")
	}

	db, err := sql.Open("sqlite3", "thyme.db")
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	rows, err := db.Query(`
	SELECT path FROM photos
	JOIN sets ON photos.set_id = sets.id
	ORDER BY sets.taken_at DESC, photos.taken_at ASC
	`)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	if err := os.MkdirAll(thumbsDir, os.ModeDir|0755); err != nil {
		log.Fatal(err)
	}

	ch := make(chan string)
	wg := sync.WaitGroup{}

	for i := 0; i < workers; i++ {
		wg.Add(1)
		go generateThumbs(ch, &wg)
	}

	for rows.Next() {
		var photoPath string
		rows.Scan(&photoPath)
		ch <- photoPath
	}

	close(ch)
	wg.Wait()

	if err := rows.Err(); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Make generateThumbsImpl return an error<commit_after>package main

import (
	"crypto/md5"
	"database/sql"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path"
	"sync"

	_ "github.com/mattn/go-sqlite3"
)

const (
	thumbsDir      = "public/thumbs"
	bigThumbSize   = "1000x1000"
	smallThumbSize = "200x200"
	workers        = 4 // min: 1
)

func generateSmallThumb(photoPath, identifier string) (thumbPath string, err error) {
	thumbPath = path.Join(thumbsDir, fmt.Sprintf("%s_small.jpg", identifier))
	if _, err = os.Stat(thumbPath); os.IsNotExist(err) { // file does not exist
		err = exec.Command(
			"convert", photoPath,
			"-auto-orient",
			"-strip",
			"-resize", smallThumbSize+"^",
			"-gravity", "center",
			"-extent", smallThumbSize,
			thumbPath).Run()
	}
	return
}

func generateBigThumb(photoPath, identifier string) (thumbPath string, err error) {
	thumbPath = path.Join(thumbsDir, fmt.Sprintf("%s_big.jpg", identifier))
	if _, err = os.Stat(thumbPath); os.IsNotExist(err) { // file does not exist
		err = exec.Command(
			"convert", photoPath,
			"-auto-orient",
			"-strip",
			"-resize", bigThumbSize,
			thumbPath).Run()
	}
	return
}

func generateThumbsImpl(photoPath string) (err error) {
	identifier := fmt.Sprintf("%x", md5.Sum([]byte(photoPath)))

	bigThumbPath, err := generateBigThumb(photoPath, identifier)
	if err != nil { // error
		return
	}

	fmt.Println(bigThumbPath)

	smallThumbPath, err := generateSmallThumb(bigThumbPath, identifier)
	if err == nil { // success
		fmt.Println(smallThumbPath)
	}

	return
}

func generateThumbs(ch chan string, wg *sync.WaitGroup) {
	defer wg.Done()

	for photoPath := range ch {
		generateThumbsImpl(photoPath)
	}
}

func main() {
	if workers < 1 {
		log.Fatal("number of workers must be at least 1")
	}

	db, err := sql.Open("sqlite3", "thyme.db")
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	rows, err := db.Query(`
	SELECT path FROM photos
	JOIN sets ON photos.set_id = sets.id
	ORDER BY sets.taken_at DESC, photos.taken_at ASC
	`)
	if err != nil {
		log.Fatal(err)
	}
	defer rows.Close()

	if err := os.MkdirAll(thumbsDir, os.ModeDir|0755); err != nil {
		log.Fatal(err)
	}

	ch := make(chan string)
	wg := sync.WaitGroup{}

	for i := 0; i < workers; i++ {
		wg.Add(1)
		go generateThumbs(ch, &wg)
	}

	for rows.Next() {
		var photoPath string
		rows.Scan(&photoPath)
		ch <- photoPath
	}

	close(ch)
	wg.Wait()

	if err := rows.Err(); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package reduction
<commit_msg>Delete polynomial.go<commit_after><|endoftext|>