<commit_before>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package mysql

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/webx-top/com"
	"github.com/webx-top/echo"

	"github.com/admpub/archiver"
	"github.com/admpub/errors"
	loga "github.com/admpub/log"
	"github.com/admpub/nging/v3/application/handler"
	"github.com/admpub/nging/v3/application/library/dbmanager/driver"
	"github.com/admpub/nging/v3/application/library/dbmanager/driver/mysql/utils"
	"github.com/admpub/nging/v3/application/library/notice"
)

// SQLTempDir sql文件缓存目录获取函数(用于导入导出SQL)
var SQLTempDir = os.TempDir

func TempDir(op string) string {
	dir := filepath.Join(SQLTempDir(), `dbmanager/cache`, op)
	err := com.MkdirAll(dir, os.ModePerm)
	if err != nil {
		loga.Error(err)
	}
	return dir
}

func (m *mySQL) getCharsetList() []string {
	e, _ := m.getCharsets()
	cs := make([]string, 0)
	for k := range e {
		cs = append(cs, k)
	}
	sort.Strings(cs)
	return cs
}

func (m *mySQL) exporting() error {
	return m.bgExecManage(utils.OpExport)
}

func (m *mySQL) bgExecManage(op utils.OP) error {
	var err error
	if m.IsPost() {
		data := m.Data()
		keys := m.FormValues(`key`)
		for _, key := range keys {
			err = utils.Cancel(op, key)
			if err != nil {
				return m.JSON(data.SetError(err))
			}
		}
		data.SetInfo(m.T(`操作成功`))
		m.ok(m.T(`操作成功`))
		return m.returnTo(m.GenURL(op.String()) + `&process=1`)
	}
	m.Set(`op`, op)
	m.Set(`list`, utils.ListBy(op))
	var title string
	if op == utils.OpExport {
		title = m.T(`导出SQL`)
	} else {
		title = m.T(`导入SQL`)
	}
	m.Set(`title`, title)
	m.Set(`cacheDir`, handler.URLFor(`/download/file?path=dbmanager/cache/`+op.String()))
	return m.Render(`db/mysql/process_store`, m.checkErr(err))
}

func (m *mySQL) Export() error {
	//fmt.Printf("%#v\n", m.getCharsetList())
	process := m.Queryx(`process`).Bool()
	if process {
		return m.exporting()
	}
	var err error
	if m.IsPost() {
		var tables []string
		if len(m.dbName) == 0 {
			m.fail(m.T(`请选择数据库`))
			return m.returnTo(m.GenURL(`listDb`))
		}
		if m.Formx(`all`).Bool() {
			var ok bool
			tables, ok = m.Get(`tableList`).([]string)
			if !ok {
				tables, err = m.getTables()
				if err != nil {
					m.fail(err.Error())
					return m.returnTo(m.GenURL(`export`))
				}
			}
		} else {
			tables = m.FormValues(`table`)
			if len(tables) == 1 && len(tables[0]) > 0 {
				tables = strings.Split(tables[0], `,`)
			}
			views := m.FormValues(`view`)
			if len(views) == 1 && len(views[0]) > 0 {
				views = append(tables, strings.Split(views[0], `,`)...)
			}
			if len(views) > 0 {
				tables = append(tables, views...)
			}
		}
		output := m.Form(`output`)
		types := m.FormValues(`type`)
		cacheKey := com.Md5(com.Dump([]interface{}{tables, output, types}, false))

		var exports utils.Exec
		if old, exists := utils.Backgrounds.Load(utils.OpExport); exists {
			exports = old.(utils.Exec)
		} else {
			exports = utils.Exec{}
		}
		if exports.Exists(cacheKey) {
			return errors.New(m.T(`任务正在后台处理中，请稍候...`))
		}

		var (
			structWriter, dataWriter interface{}
			sqlFiles                 []string
			async                    bool
			bgExec                   = utils.NewGBExec(context.TODO(), echo.H{
				`database`: m.dbName,
				`tables`:   tables,
				`output`:   output,
				`types`:    types,
			})
			fileInfos = bgExec.Procs
		)
		exports[cacheKey] = bgExec
		nowTime := com.String(time.Now().Unix())
		saveDir := TempDir(`export`)
		switch output {
		case `down`:
			m.Response().Header().Set(echo.HeaderContentType, echo.MIMEOctetStream)
			m.Response().Header().Set(echo.HeaderContentDisposition, fmt.Sprintf("attachment; filename=%q", m.dbName+"-sql-"+nowTime+".sql"))
			fallthrough
		case `open`:
			m.Response().Header().Set(echo.HeaderContentType, echo.MIMETextPlainCharsetUTF8)
			if com.InSlice(`struct`, types) {
				structWriter = m.Response()
			}
			if com.InSlice(`data`, types) {
				dataWriter = m.Response()
			}
		default:
			async = true
			if com.InSlice(`struct`, types) {
				structFile := filepath.Join(saveDir, m.dbName+`-struct-`+nowTime+`.sql`)
				sqlFiles = append(sqlFiles, structFile)
				structWriter = structFile
				fi := &utils.FileInfo{
					Start: time.Now(),
					Path:  structFile,
				}
				*fileInfos = append(*fileInfos, fi)
			}
			if com.InSlice(`data`, types) {
				dataFile := filepath.Join(saveDir, m.dbName+`-data-`+nowTime+`.sql`)
				sqlFiles = append(sqlFiles, dataFile)
				dataWriter = dataFile
				fi := &utils.FileInfo{
					Start: time.Now(),
					Path:  dataFile,
				}
				*fileInfos = append(*fileInfos, fi)
			}
		}
		cfg := *m.DbAuth
		cfg.Db = m.dbName
		coll, err := m.getCollation(m.dbName, nil)
		if err != nil {
			return err
		}
		cfg.Charset = strings.SplitN(coll, `_`, 2)[0]

		user := handler.User(m.Context)
		var username string
		if user != nil {
			username = user.Username
		}
		noticer := notice.New(m.Context, `databaseExport`, username)

		worker := func(ctx context.Context, cfg driver.DbAuth) error {
			defer func() {
				if r := recover(); r != nil {
					err = fmt.Errorf(`RECOVER: %v`, r)
				}
			}()
			err = utils.Export(ctx, noticer, &cfg, tables, structWriter, dataWriter, true)
			if err != nil {
				loga.Error(err)
				return err
			}
			if len(sqlFiles) > 0 {
				now := time.Now()
				for _, fi := range *fileInfos {
					fi.End = now
					fi.Size, err = com.FileSize(fi.Path)
					if err != nil {
						fi.Error = err.Error()
					}
					fi.Elapsed = fi.End.Sub(fi.Start)
				}
				zipFile := filepath.Join(saveDir, cfg.Db+"-sql-"+nowTime+".zip")
				fi := &utils.FileInfo{
					Start:      now,
					Path:       zipFile,
					Compressed: true,
				}
				err = archiver.Archive(sqlFiles, zipFile)
				if err != nil {
					loga.Error(err)
					return err
				}
				for _, sqlFile := range sqlFiles {
					os.Remove(sqlFile)
				}
				fi.Size, err = com.FileSize(zipFile)
				if err != nil {
					fi.Error = err.Error()
				}
				fi.End = time.Now()
				fi.Elapsed = fi.End.Sub(fi.Start)
				*fileInfos = append(*fileInfos, fi)
				ioutil.WriteFile(zipFile+`.txt`, com.Str2bytes(com.Dump(fileInfos, false)), os.ModePerm)
				exports.Cancel(cacheKey)
			}
			return nil
		}
		if !async {
			done := make(chan struct{})
			ctx := m.Request().StdRequest().Context()
			go func() {
				for {
					select {
					case <-ctx.Done():
						bgExec.Cancel()()
						return
					case <-done:
						return
					}
				}
			}()
			err = worker(bgExec.Context(), cfg)
			if err != nil {
				noticer(m.T(`导出失败`)+`: `+err.Error(), 0)
			}
			done <- struct{}{}
			return err
		}
		data := m.Data()
		data.SetInfo(m.T(`任务已经在后台成功启动`))
		data.SetURL(handler.URLFor(`/download/file?path=dbmanager/cache/export`))
		go worker(m, cfg)
		exports.Add(utils.OpExport, cacheKey, bgExec)
		return m.JSON(data)
	}
	return m.Redirect(m.GenURL(`listTable`, m.dbName))
}
<commit_msg>improved<commit_after>/*
   Nging is a toolbox for webmasters
   Copyright (C) 2018-present  Wenhui Shen <swh@admpub.com>

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published
   by the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package mysql

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/webx-top/com"
	"github.com/webx-top/echo"

	"github.com/admpub/archiver"
	"github.com/admpub/errors"
	loga "github.com/admpub/log"
	"github.com/admpub/nging/v3/application/handler"
	"github.com/admpub/nging/v3/application/library/dbmanager/driver"
	"github.com/admpub/nging/v3/application/library/dbmanager/driver/mysql/utils"
	"github.com/admpub/nging/v3/application/library/notice"
)

// SQLTempDir sql文件缓存目录获取函数(用于导入导出SQL)
var SQLTempDir = os.TempDir

func TempDir(op string) string {
	dir := filepath.Join(SQLTempDir(), `dbmanager/cache`, op)
	err := com.MkdirAll(dir, os.ModePerm)
	if err != nil {
		loga.Error(err)
	}
	return dir
}

func (m *mySQL) getCharsetList() []string {
	e, _ := m.getCharsets()
	cs := make([]string, 0)
	for k := range e {
		cs = append(cs, k)
	}
	sort.Strings(cs)
	return cs
}

func (m *mySQL) exporting() error {
	return m.bgExecManage(utils.OpExport)
}

func (m *mySQL) bgExecManage(op utils.OP) error {
	var err error
	if m.IsPost() {
		data := m.Data()
		keys := m.FormValues(`key`)
		for _, key := range keys {
			err = utils.Cancel(op, key)
			if err != nil {
				return m.JSON(data.SetError(err))
			}
		}
		data.SetInfo(m.T(`操作成功`))
		m.ok(m.T(`操作成功`))
		return m.returnTo(m.GenURL(op.String()) + `&process=1`)
	}
	m.Set(`op`, op)
	m.Set(`list`, utils.ListBy(op))
	var title string
	if op == utils.OpExport {
		title = m.T(`导出SQL`)
	} else {
		title = m.T(`导入SQL`)
	}
	m.Set(`title`, title)
	m.Set(`cacheDir`, handler.URLFor(`/download/file?path=dbmanager/cache/`+op.String()))
	return m.Render(`db/mysql/process_store`, m.checkErr(err))
}

func (m *mySQL) Export() error {
	//fmt.Printf("%#v\n", m.getCharsetList())
	process := m.Queryx(`process`).Bool()
	if process {
		return m.exporting()
	}
	var err error
	if m.IsPost() {
		var tables []string
		if len(m.dbName) == 0 {
			m.fail(m.T(`请选择数据库`))
			return m.returnTo(m.GenURL(`listDb`))
		}
		if m.Formx(`all`).Bool() {
			var ok bool
			tables, ok = m.Get(`tableList`).([]string)
			if !ok {
				tables, err = m.getTables()
				if err != nil {
					m.fail(err.Error())
					return m.returnTo(m.GenURL(`export`))
				}
			}
		} else {
			tables = m.FormValues(`table`)
			if len(tables) == 1 && len(tables[0]) > 0 {
				tables = strings.Split(tables[0], `,`)
			}
			views := m.FormValues(`view`)
			if len(views) == 1 && len(views[0]) > 0 {
				views = append(tables, strings.Split(views[0], `,`)...)
			}
			if len(views) > 0 {
				tables = append(tables, views...)
			}
		}
		output := m.Form(`output`)
		types := m.FormValues(`type`)
		cacheKey := com.Md5(com.Dump([]interface{}{tables, output, types}, false))

		var exports utils.Exec
		if old, exists := utils.Backgrounds.Load(utils.OpExport); exists {
			exports = old.(utils.Exec)
		} else {
			exports = utils.Exec{}
		}
		if exports.Exists(cacheKey) {
			return errors.New(m.T(`任务正在后台处理中，请稍候...`))
		}

		var (
			structWriter, dataWriter interface{}
			sqlFiles                 []string
			dbSaveDir                string
			async                    bool
			bgExec                   = utils.NewGBExec(context.TODO(), echo.H{
				`database`: m.dbName,
				`tables`:   tables,
				`output`:   output,
				`types`:    types,
			})
			fileInfos = bgExec.Procs
		)
		exports[cacheKey] = bgExec
		nowTime := time.Now().Local().Format("20060102150405.000")
		saveDir := TempDir(`export`)
		switch output {
		case `down`:
			m.Response().Header().Set(echo.HeaderContentType, echo.MIMEOctetStream)
			m.Response().Header().Set(echo.HeaderContentDisposition, fmt.Sprintf("attachment; filename=%q", m.dbName+"-sql-"+nowTime+".sql"))
			fallthrough
		case `open`:
			m.Response().Header().Set(echo.HeaderContentType, echo.MIMETextPlainCharsetUTF8)
			if com.InSlice(`struct`, types) {
				structWriter = m.Response()
			}
			if com.InSlice(`data`, types) {
				dataWriter = m.Response()
			}
		default:
			async = true
			dbSaveDir = filepath.Join(saveDir, m.dbName)
			com.MkdirAll(dbSaveDir, os.ModePerm)
			if com.InSlice(`struct`, types) {
				structFile := filepath.Join(dbSaveDir, `struct-`+nowTime+`.sql`)
				sqlFiles = append(sqlFiles, structFile)
				structWriter = structFile
				fi := &utils.FileInfo{
					Start: time.Now(),
					Path:  structFile,
				}
				*fileInfos = append(*fileInfos, fi)
			}
			if com.InSlice(`data`, types) {
				dataFile := filepath.Join(dbSaveDir, `data-`+nowTime+`.sql`)
				sqlFiles = append(sqlFiles, dataFile)
				dataWriter = dataFile
				fi := &utils.FileInfo{
					Start: time.Now(),
					Path:  dataFile,
				}
				*fileInfos = append(*fileInfos, fi)
			}
		}
		cfg := *m.DbAuth
		cfg.Db = m.dbName
		coll, err := m.getCollation(m.dbName, nil)
		if err != nil {
			return err
		}
		cfg.Charset = strings.SplitN(coll, `_`, 2)[0]

		user := handler.User(m.Context)
		var username string
		if user != nil {
			username = user.Username
		}
		noticer := notice.New(m.Context, `databaseExport`, username)

		worker := func(ctx context.Context, cfg driver.DbAuth) error {
			defer func() {
				if r := recover(); r != nil {
					err = fmt.Errorf(`RECOVER: %v`, r)
				}
			}()
			err = utils.Export(ctx, noticer, &cfg, tables, structWriter, dataWriter, true)
			if err != nil {
				loga.Error(err)
				return err
			}
			if len(sqlFiles) > 0 {
				now := time.Now()
				for _, fi := range *fileInfos {
					fi.End = now
					fi.Size, err = com.FileSize(fi.Path)
					if err != nil {
						fi.Error = err.Error()
					}
					fi.Elapsed = fi.End.Sub(fi.Start)
				}
				zipFile := filepath.Join(dbSaveDir, "sql-"+nowTime+".zip")
				fi := &utils.FileInfo{
					Start:      now,
					Path:       zipFile,
					Compressed: true,
				}
				err = archiver.Archive(sqlFiles, zipFile)
				if err != nil {
					loga.Error(err)
					return err
				}
				for _, sqlFile := range sqlFiles {
					os.Remove(sqlFile)
				}
				fi.Size, err = com.FileSize(zipFile)
				if err != nil {
					fi.Error = err.Error()
				}
				fi.End = time.Now()
				fi.Elapsed = fi.End.Sub(fi.Start)
				*fileInfos = append(*fileInfos, fi)
				ioutil.WriteFile(zipFile+`.txt`, com.Str2bytes(com.Dump(fileInfos, false)), os.ModePerm)
				exports.Cancel(cacheKey)
			}
			return nil
		}
		if !async {
			done := make(chan struct{})
			ctx := m.Request().StdRequest().Context()
			go func() {
				for {
					select {
					case <-ctx.Done():
						bgExec.Cancel()()
						return
					case <-done:
						return
					}
				}
			}()
			err = worker(bgExec.Context(), cfg)
			if err != nil {
				noticer(m.T(`导出失败`)+`: `+err.Error(), 0)
			}
			done <- struct{}{}
			return err
		}
		data := m.Data()
		data.SetInfo(m.T(`任务已经在后台成功启动`))
		data.SetURL(handler.URLFor(`/download/file?path=dbmanager/cache/export`))
		go worker(m, cfg)
		exports.Add(utils.OpExport, cacheKey, bgExec)
		return m.JSON(data)
	}
	return m.Redirect(m.GenURL(`listTable`, m.dbName))
}
<|endoftext|>


<commit_before>// quick sorting

package main

import (
	"fmt"
	"ashumeow/meow_sort"
)

func quick(meow []int) []int {
	if len(meow) <= 1 {
		return meow
	}
	median := meow[rand.Intn(len(meow))]
	low := make([]int,0,len(meow))
	high := make([]int,0,len(meow))
	middle := make([]int,0,len(meow))

	for ig_ig, stuffs := range meow {
		switch {
			case stuffs < median:
				low = append(low, stuffs)
			case stuffs == median:
				middle = append(middle, stuffs)
			case stuffs > median:
				high = append(high, stuffs)
		}
	}
	low = quick(low)
	high = quick(high)
	low = append(low, middle)
	low = append(low, high)
	return low
}

func main() {
	meow := meow_sort.RandArray(10)
	fmt.Println("Given array is: ",meow)
	fmt.Println("")
	fmt.Println("Sorted Array is: ",quick(meow))
}<commit_msg>update<commit_after>// quick sorting

package main

import (
	"fmt"
	"math/rand"
	"ashumeow/meow_sort"
)

func quick(meow []int) []int {
	if len(meow) <= 1 {
		return meow
	}
	median := meow[rand.Intn(len(meow))]
	low := make([]int,0,len(meow))
	high := make([]int,0,len(meow))
	middle := make([]int,0,len(meow))

	for ig_ig, stuffs := range meow {
		switch {
			case stuffs < median:
				low = append(low, stuffs)
			case stuffs == median:
				middle = append(middle, stuffs)
			case stuffs > median:
				high = append(high, stuffs)
		}
	}
	low = quick(low)
	high = quick(high)
	low = append(low, middle)
	low = append(low, high)
	return low
}

func main() {
	meow := meow_sort.RandArray(10)
	fmt.Println("Given array is: ",meow)
	fmt.Println("")
	fmt.Println("Sorted Array is: ",quick(meow))
}<|endoftext|>


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

package custom

import (
	"bufio"
	"bytes"
	"encoding"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"reflect"
	"time"
	"unsafe"

	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/timestamppb"
)

var (
	_ = json.Marshaler((*DateTime)(nil))
	_ = json.Unmarshaler((*DateTime)(nil))
	_ = encoding.TextMarshaler((*DateTime)(nil))
	_ = encoding.TextUnmarshaler((*DateTime)(nil))
	_ = xml.Marshaler((*DateTime)(nil))
	_ = xml.Unmarshaler((*DateTime)(nil))
)

type DateTime struct {
	Time time.Time
}

func getWriter(enc *xml.Encoder) *bufio.Writer {
	rEnc := reflect.ValueOf(enc)
	rP := rEnc.Elem().FieldByName("p").Addr()
	return *(**bufio.Writer)(unsafe.Pointer(rP.Elem().FieldByName("Writer").UnsafeAddr()))
}

func (dt *DateTime) MarshalXML(enc *xml.Encoder, start xml.StartElement) error {
	if dt.IsZero() {
		start.Attr = append(start.Attr,
			xml.Attr{Name: xml.Name{Space: "http://www.w3.org/2001/XMLSchema-instance", Local: "nil"}, Value: "true"})

		bw := getWriter(enc)
		bw.Flush()
		old := *bw
		var buf bytes.Buffer
		*bw = *bufio.NewWriter(&buf)
		if err := enc.EncodeElement("", start); err != nil {
			return err
		}
		b := bytes.ReplaceAll(bytes.ReplaceAll(bytes.ReplaceAll(bytes.ReplaceAll(
		buf.Bytes(),
			[]byte("_XMLSchema-instance:"), []byte("xsi:")),
			[]byte("xmlns:_XMLSchema-instance="), []byte("xmlns:xsi=")),
			[]byte("XMLSchema-instance:"), []byte("xsi:")),
			[]byte("xmlns:XMLSchema-instance="), []byte("xmlns:xsi="))
		*bw = old
		bw.Write(b)
		return bw.Flush()
	}
	return enc.EncodeElement(dt.Time.In(time.Local).Format(time.RFC3339), start)
}
func (dt *DateTime) UnmarshalXML(dec *xml.Decoder, st xml.StartElement) error {
	var s string
	if err := dec.DecodeElement(&s, &st); err != nil {
		return err
	}
	return dt.UnmarshalText([]byte(s))
}

func (dt *DateTime) IsZero() (zero bool) {
	//defer func() { log.Printf("IsZero(%#v): %t", dt, zero) }()
	if dt == nil {
		return true
	}
	defer func() {
		if r := recover(); r != nil {
			zero = true
		}
	}()
	return dt.Time.IsZero()
}
func (dt *DateTime) MarshalJSON() ([]byte, error) {
	if dt.IsZero() {
		return []byte(`""`), nil
	}
	return dt.Time.In(time.Local).MarshalJSON()
}
func (dt *DateTime) UnmarshalJSON(data []byte) error {
	// Ignore null, like in the main JSON package.
	data = bytes.TrimSpace(data)
	if len(data) == 0 || bytes.Equal(data, []byte(`""`)) || bytes.Equal(data, []byte("null")) {
		return nil
	}
	return dt.UnmarshalText(data)
}

// MarshalText implements the encoding.TextMarshaler interface.
// The time is formatted in RFC 3339 format, with sub-second precision added if present.
func (dt *DateTime) MarshalText() ([]byte, error) {
	if dt.IsZero() {
		return nil, nil
	}
	return dt.Time.In(time.Local).MarshalText()
}

// UnmarshalText implements the encoding.TextUnmarshaler interface.
// The time is expected to be in RFC 3339 format.
func (dt *DateTime) UnmarshalText(data []byte) error {
	data = bytes.Trim(data, " \"")
	n := len(data)
	if n == 0 {
		dt.Time = time.Time{}
		//log.Println("time=")
		return nil
	}
	if n > len(time.RFC3339) {
		n = len(time.RFC3339)
	} else if n < 4 {
		n = 4
	} else if n > 10 && data[10] != time.RFC3339[10] {
		data[10] = time.RFC3339[10]
	}
	var err error
	// Fractional seconds are handled implicitly by Parse.
	dt.Time, err = time.ParseInLocation(time.RFC3339[:n], string(data), time.Local)
	//log.Printf("s=%q time=%v err=%+v", data, dt.Time, err)
	if err != nil {
		return fmt.Errorf("%s: %w", string(data), err)
	}
	return nil
}

func (dt *DateTime) Timestamp() *timestamppb.Timestamp {
	if dt.IsZero() {
		return nil
	}
	return timestamppb.New(dt.Time)
}
func (dt *DateTime) MarshalTo(dAtA []byte) (int, error) {
	if dt.IsZero() {
		return 0, nil
	}
	b, err := proto.MarshalOptions{}.MarshalAppend(dAtA[:0], dt.Timestamp())
	_ = dAtA[len(b)-1] // panic if buffer is too short
	return len(b), err
}
func (dt *DateTime) Marshal() (dAtA []byte, err error) {
	if dt.IsZero() {
		return nil, nil
	}
	return proto.Marshal(dt.Timestamp())
}
func (dt *DateTime) String() string {
	if dt.IsZero() {
		return ""
	}
	return dt.Time.In(time.Local).Format(time.RFC3339)
}

func (dt *DateTime) ProtoMessage() {}

func (dt *DateTime) ProtoSize() (n int) {
	if dt.IsZero() {
		return 0
	}
	return proto.Size(dt.Timestamp())
}
func (dt *DateTime) Reset() {
	if dt != nil {
		dt.Time = time.Time{}
	}
}
func (dt *DateTime) Size() (n int) {
	if dt.IsZero() {
		return 0
	}
	return proto.Size(dt.Timestamp())
}
func (dt *DateTime) Unmarshal(dAtA []byte) error {
	var ts timestamppb.Timestamp
	if err := proto.Unmarshal(dAtA, &ts); err != nil {
		return err
	}
	if ts.Seconds == 0 && ts.Nanos == 0 {
		dt.Time = time.Time{}
	} else {
		dt.Time = ts.AsTime()
	}
	return nil
}
<commit_msg>custom.DateTime: Define Marshal* on DateTime, not *DateTime<commit_after>// Copyright 2019, 2020 Tamás Gulácsi
//
//
//    Licensed under the Apache License, Version 2.0 (the "License");
//    you may not use this file except in compliance with the License.
//    You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//    Unless required by applicable law or agreed to in writing, software
//    distributed under the License is distributed on an "AS IS" BASIS,
//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//    See the License for the specific language governing permissions and
//    limitations under the License.

package custom

import (
	"bufio"
	"bytes"
	"encoding"
	"encoding/json"
	"encoding/xml"
	"fmt"
	"reflect"
	"time"
	"unsafe"

	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/types/known/timestamppb"
)

var (
	_ = json.Marshaler((*DateTime)(nil))
	_ = json.Unmarshaler((*DateTime)(nil))
	_ = encoding.TextMarshaler((*DateTime)(nil))
	_ = encoding.TextUnmarshaler((*DateTime)(nil))
	_ = xml.Marshaler((*DateTime)(nil))
	_ = xml.Unmarshaler((*DateTime)(nil))
)

type DateTime struct {
	Time time.Time
}

func getWriter(enc *xml.Encoder) *bufio.Writer {
	rEnc := reflect.ValueOf(enc)
	rP := rEnc.Elem().FieldByName("p").Addr()
	return *(**bufio.Writer)(unsafe.Pointer(rP.Elem().FieldByName("Writer").UnsafeAddr()))
}

func (dt DateTime) MarshalXML(enc *xml.Encoder, start xml.StartElement) error {
	if dt.IsZero() {
		start.Attr = append(start.Attr,
			xml.Attr{Name: xml.Name{Space: "http://www.w3.org/2001/XMLSchema-instance", Local: "nil"}, Value: "true"})

		bw := getWriter(enc)
		bw.Flush()
		old := *bw
		var buf bytes.Buffer
		*bw = *bufio.NewWriter(&buf)
		if err := enc.EncodeElement("", start); err != nil {
			return err
		}
		b := bytes.ReplaceAll(bytes.ReplaceAll(bytes.ReplaceAll(bytes.ReplaceAll(
			buf.Bytes(),
			[]byte("_XMLSchema-instance:"), []byte("xsi:")),
			[]byte("xmlns:_XMLSchema-instance="), []byte("xmlns:xsi=")),
			[]byte("XMLSchema-instance:"), []byte("xsi:")),
			[]byte("xmlns:XMLSchema-instance="), []byte("xmlns:xsi="))
		*bw = old
		bw.Write(b)
		return bw.Flush()
	}
	return enc.EncodeElement(dt.Time.In(time.Local).Format(time.RFC3339), start)
}
func (dt *DateTime) UnmarshalXML(dec *xml.Decoder, st xml.StartElement) error {
	var s string
	if err := dec.DecodeElement(&s, &st); err != nil {
		return err
	}
	return dt.UnmarshalText([]byte(s))
}

func (dt *DateTime) IsZero() (zero bool) {
	//defer func() { log.Printf("IsZero(%#v): %t", dt, zero) }()
	if dt == nil {
		return true
	}
	defer func() {
		if r := recover(); r != nil {
			zero = true
		}
	}()
	return dt.Time.IsZero()
}
func (dt DateTime) MarshalJSON() ([]byte, error) {
	if dt.IsZero() {
		return []byte(`""`), nil
	}
	return dt.Time.In(time.Local).MarshalJSON()
}
func (dt *DateTime) UnmarshalJSON(data []byte) error {
	// Ignore null, like in the main JSON package.
	data = bytes.TrimSpace(data)
	if len(data) == 0 || bytes.Equal(data, []byte(`""`)) || bytes.Equal(data, []byte("null")) {
		dt.Time = time.Time{}
		return nil
	}
	return dt.UnmarshalText(data)
}

// MarshalText implements the encoding.TextMarshaler interface.
// The time is formatted in RFC 3339 format, with sub-second precision added if present.
func (dt DateTime) MarshalText() ([]byte, error) {
	if dt.IsZero() {
		return nil, nil
	}
	return dt.Time.In(time.Local).MarshalText()
}

// UnmarshalText implements the encoding.TextUnmarshaler interface.
// The time is expected to be in RFC 3339 format.
func (dt *DateTime) UnmarshalText(data []byte) error {
	data = bytes.Trim(data, " \"")
	n := len(data)
	if n == 0 {
		dt.Time = time.Time{}
		//log.Println("time=")
		return nil
	}
	if n > len(time.RFC3339) {
		n = len(time.RFC3339)
	} else if n < 4 {
		n = 4
	} else if n > 10 && data[10] != time.RFC3339[10] {
		data[10] = time.RFC3339[10]
	}
	var err error
	// Fractional seconds are handled implicitly by Parse.
	dt.Time, err = time.ParseInLocation(time.RFC3339[:n], string(data), time.Local)
	//log.Printf("s=%q time=%v err=%+v", data, dt.Time, err)
	if err != nil {
		return fmt.Errorf("ParseInLocation(%q, %q): %w", time.RFC3339[:n], string(data), err)
	}
	return nil
}

func (dt *DateTime) Timestamp() *timestamppb.Timestamp {
	if dt.IsZero() {
		return nil
	}
	return timestamppb.New(dt.Time)
}
func (dt *DateTime) MarshalTo(dAtA []byte) (int, error) {
	if dt.IsZero() {
		return 0, nil
	}
	b, err := proto.MarshalOptions{}.MarshalAppend(dAtA[:0], dt.Timestamp())
	_ = dAtA[len(b)-1] // panic if buffer is too short
	return len(b), err
}
func (dt *DateTime) Marshal() (dAtA []byte, err error) {
	if dt.IsZero() {
		return nil, nil
	}
	return proto.Marshal(dt.Timestamp())
}
func (dt *DateTime) String() string {
	if dt.IsZero() {
		return ""
	}
	return dt.Time.In(time.Local).Format(time.RFC3339)
}

func (dt *DateTime) ProtoMessage() {}

func (dt *DateTime) ProtoSize() (n int) {
	if dt.IsZero() {
		return 0
	}
	return proto.Size(dt.Timestamp())
}
func (dt *DateTime) Reset() {
	if dt != nil {
		dt.Time = time.Time{}
	}
}
func (dt *DateTime) Size() (n int) {
	if dt.IsZero() {
		return 0
	}
	return proto.Size(dt.Timestamp())
}
func (dt *DateTime) Unmarshal(dAtA []byte) error {
	var ts timestamppb.Timestamp
	if err := proto.Unmarshal(dAtA, &ts); err != nil {
		return err
	}
	if ts.Seconds == 0 && ts.Nanos == 0 {
		dt.Time = time.Time{}
	} else {
		dt.Time = ts.AsTime()
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package features

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"k8s.io/contrib/mungegithub/github"
	"k8s.io/kubernetes/pkg/util/sets"
	"k8s.io/kubernetes/pkg/util/yaml"

	parseYaml "github.com/ghodss/yaml"
	"github.com/golang/glog"
	"github.com/spf13/cobra"
)

const (
	ownerFilename = "OWNERS" // file which contains owners and assignees
	// RepoFeatureName is how mungers should indicate this is required
	RepoFeatureName = "gitrepos"
)

type assignmentConfig struct {
	Assignees []string `json:assignees yaml:assignees`
	Approvers []string `json:approvers yaml:approvers`
	//Reviewers []string `json:reviewers yaml:reviewers`
	//Owners []string `json:owners`
}

// RepoInfo provides information about users in OWNERS files in a git repo
type RepoInfo struct {
	BaseDir      string
	EnableMdYaml bool

	enabled    bool
	projectDir string
	assignees  map[string]sets.String
	config     *github.Config
	//owners   map[string]sets.String
}

func init() {
	RegisterFeature(&RepoInfo{})
}

// Name is just going to return the name mungers use to request this feature
func (o *RepoInfo) Name() string {
	return RepoFeatureName
}

func (o *RepoInfo) walkFunc(path string, info os.FileInfo, err error) error {
	if err != nil {
		glog.Errorf("%v", err)
		return nil
	}
	filename := filepath.Base(path)
	if info.Mode().IsDir() {
		switch filename {
		case ".git":
			return filepath.SkipDir
		case "_output":
			return filepath.SkipDir
		}
	}
	if !info.Mode().IsRegular() {
		return nil
	}

	c := &assignmentConfig{}

	// '.md' files may contain assignees at the top of the file in a yaml header
	// Flag guarded because this is only enabled in some repos
	if o.EnableMdYaml && filename != ownerFilename && strings.HasSuffix(filename, "md") {
		// Parse the yaml header from the file if it exists and marshal into the config
		if err := decodeAssignmentConfig(path, c); err != nil {
			glog.Errorf("%v", err)
			return err
		}

		// Set assignees for this file using the relative path if they were found
		path, err = filepath.Rel(o.projectDir, path)
		if err != nil {
			glog.Errorf("Unable to find relative path between %q and %q: %v", o.projectDir, path, err)
			return err
		}
		if len(c.Assignees) > 0 {
			o.assignees[path] = sets.NewString(c.Assignees...)
		}
		return nil
	}

	if filename != ownerFilename {
		return nil
	}

	file, err := os.Open(path)
	if err != nil {
		glog.Errorf("%v", err)
		return nil
	}
	defer file.Close()

	if err := yaml.NewYAMLToJSONDecoder(file).Decode(c); err != nil {
		glog.Errorf("%v", err)
		return nil
	}

	path, err = filepath.Rel(o.projectDir, path)
	if err != nil {
		glog.Errorf("Unable to find relative path between %q and %q: %v", o.projectDir, path, err)
		return err
	}
	path = filepath.Dir(path)
	// Make the root explicitly / so its easy to distinguish. Nothing else is `/` anchored
	if path == "." {
		path = "/"
	}
	if len(c.Assignees) > 0 {
		o.assignees[path] = sets.NewString(c.Assignees...)
	} else if len(c.Approvers) > 0 {
		o.assignees[path] = sets.NewString(c.Approvers...)
	}
	//if len(c.Owners) > 0 {
	//o.owners[path] = sets.NewString(c.Owners...)
	//}
	return nil
}

// decodeAssignmentConfig will parse the yaml header if it exists and unmarshal it into an assignmentConfig.
// If no yaml header is found, do nothing
// Returns an error if the file cannot be read or the yaml header is found but cannot be unmarshalled
var mdStructuredHeaderRegex = regexp.MustCompile("^---\n(.|\n)*\n---")

func decodeAssignmentConfig(path string, config *assignmentConfig) error {
	fileBytes, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}
	// Parse the yaml header from the top of the file.  Will return an empty string if regex does not match.
	meta := mdStructuredHeaderRegex.FindString(string(fileBytes))

	// Unmarshal the yaml header into the config
	return parseYaml.Unmarshal([]byte(meta), &config)
}

func (o *RepoInfo) updateRepoUsers() error {
	out, err := o.GitCommand([]string{"pull"})
	if err != nil {
		glog.Errorf("Unable to run git pull:\n%s\n%v", string(out), err)
		return err
	}

	out, err = o.GitCommand([]string{"rev-parse", "HEAD"})
	if err != nil {
		glog.Errorf("Unable get sha of HEAD:\n%s\n%v", string(out), err)
		return err
	}
	sha := out

	o.assignees = map[string]sets.String{}
	//o.owners = map[string]sets.String{}
	err = filepath.Walk(o.projectDir, o.walkFunc)
	if err != nil {
		glog.Errorf("Got error %v", err)
	}
	glog.Infof("Loaded config from %s:%s", o.projectDir, sha)
	glog.V(5).Infof("assignees: %v", o.assignees)
	//glog.V(5).Infof("owners: %v", o.owners)
	return nil
}

// Initialize will initialize the munger
func (o *RepoInfo) Initialize(config *github.Config) error {
	o.enabled = true
	o.config = config
	o.projectDir = path.Join(o.BaseDir, o.config.Project)

	if len(o.BaseDir) == 0 {
		glog.Fatalf("--repo-dir is required with selected munger(s)")
	}
	finfo, err := os.Stat(o.BaseDir)
	if err != nil {
		return fmt.Errorf("Unable to stat --repo-dir: %v", err)
	}
	if !finfo.IsDir() {
		return fmt.Errorf("--repo-dir is not a directory")
	}

	// check if the cloned dir already exists, if yes, cleanup.
	if _, err := os.Stat(o.projectDir); !os.IsNotExist(err) {
		if err := o.cleanUp(o.projectDir); err != nil {
			return fmt.Errorf("Unable to remove old clone directory at %v: %v", o.projectDir, err)
		}
	}

	if cloneUrl, err := o.cloneRepo(); err != nil {
		return fmt.Errorf("Unable to clone %v: %v", cloneUrl, err)
	}
	return o.updateRepoUsers()
}

func (o *RepoInfo) cleanUp(path string) error {
	return os.RemoveAll(path)
}

func (o *RepoInfo) cloneRepo() (string, error) {
	cloneUrl := fmt.Sprintf("https://github.com/%s/%s.git", o.config.Org, o.config.Project)
	output, err := o.gitCommandDir([]string{"clone", cloneUrl, o.projectDir}, o.BaseDir)
	if err != nil {
		glog.Errorf("Failed to clone github repo: %s", output)
	}
	return cloneUrl, err
}

// EachLoop is called at the start of every munge loop
func (o *RepoInfo) EachLoop() error {
	if !o.enabled {
		return nil
	}
	_, err := o.GitCommand([]string{"remote", "update"})
	if err != nil {
		glog.Errorf("Unable to git remote update: %v", err)
	}
	return o.updateRepoUsers()
}

// AddFlags will add any request flags to the cobra `cmd`
func (o *RepoInfo) AddFlags(cmd *cobra.Command) {
	cmd.Flags().StringVar(&o.BaseDir, "repo-dir", "", "Path to perform checkout of repository")
	cmd.Flags().BoolVar(&o.EnableMdYaml, "enable-md-yaml", false, "If true, look for assignees in md yaml headers.")
}

// GitCommand will execute the git command with the `args` within the project directory.
func (o *RepoInfo) GitCommand(args []string) ([]byte, error) {
	return o.gitCommandDir(args, o.projectDir)
}

// GitCommandDir will execute the git command with the `args` within the 'dir' directory.
func (o *RepoInfo) gitCommandDir(args []string, cmdDir string) ([]byte, error) {
	cmd := exec.Command("git", args...)
	cmd.Dir = cmdDir
	return cmd.CombinedOutput()
}

func peopleForPath(path string, people map[string]sets.String, leafOnly bool, enableMdYaml bool) sets.String {
	d := path
	if !enableMdYaml {
		d = filepath.Dir(path)
	}

	out := sets.NewString()
	for {
		// special case the root
		if d == "" {
			d = "/"
		}
		s, ok := people[d]
		if ok {
			out = out.Union(s)
			if leafOnly {
				break
			}
		}
		if d == "/" {
			break
		}
		d, _ = filepath.Split(d)
		d = strings.TrimSuffix(d, "/")
	}
	return out
}

// LeafAssignees returns a set of users who are the closest assignees to the
// requested file. If pkg/OWNERS has user1 and pkg/util/OWNERS has user2 this
// will only return user2 for the path pkg/util/sets/file.go
func (o *RepoInfo) LeafAssignees(path string) sets.String {
	return peopleForPath(path, o.assignees, true, o.EnableMdYaml)
}

// Assignees returns a set of users who are the closest assignees to the
// requested file. If pkg/OWNERS has user1 and pkg/util/OWNERS has user2 this
// will return both user1 and user2 for the path pkg/util/sets/file.go
func (o *RepoInfo) Assignees(path string) sets.String {
	return peopleForPath(path, o.assignees, false, o.EnableMdYaml)
}

//func (o *RepoInfo) LeafOwners(path string) sets.String {
//return people(path, o.owners, true)
//}

//func (o *RepoInfo) Owners(path string) sets.String {
//return people(path, o.owners, false)
//}
<commit_msg>take the union of the assignees and approvers and handler approvers in md files<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package features

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"k8s.io/contrib/mungegithub/github"
	"k8s.io/kubernetes/pkg/util/sets"
	"k8s.io/kubernetes/pkg/util/yaml"

	parseYaml "github.com/ghodss/yaml"
	"github.com/golang/glog"
	"github.com/spf13/cobra"
)

const (
	ownerFilename = "OWNERS" // file which contains owners and assignees
	// RepoFeatureName is how mungers should indicate this is required
	RepoFeatureName = "gitrepos"
)

type assignmentConfig struct {
	Assignees []string `json:assignees yaml:assignees`
	Approvers []string `json:approvers yaml:approvers`
	//Reviewers []string `json:reviewers yaml:reviewers`
	//Owners []string `json:owners`
}

// RepoInfo provides information about users in OWNERS files in a git repo
type RepoInfo struct {
	BaseDir      string
	EnableMdYaml bool

	enabled    bool
	projectDir string
	assignees  map[string]sets.String
	config     *github.Config
	//owners   map[string]sets.String
}

func init() {
	RegisterFeature(&RepoInfo{})
}

// Name is just going to return the name mungers use to request this feature
func (o *RepoInfo) Name() string {
	return RepoFeatureName
}

func (o *RepoInfo) walkFunc(path string, info os.FileInfo, err error) error {
	if err != nil {
		glog.Errorf("%v", err)
		return nil
	}
	filename := filepath.Base(path)
	if info.Mode().IsDir() {
		switch filename {
		case ".git":
			return filepath.SkipDir
		case "_output":
			return filepath.SkipDir
		}
	}
	if !info.Mode().IsRegular() {
		return nil
	}

	c := &assignmentConfig{}

	// '.md' files may contain assignees at the top of the file in a yaml header
	// Flag guarded because this is only enabled in some repos
	if o.EnableMdYaml && filename != ownerFilename && strings.HasSuffix(filename, "md") {
		// Parse the yaml header from the file if it exists and marshal into the config
		if err := decodeAssignmentConfig(path, c); err != nil {
			glog.Errorf("%v", err)
			return err
		}

		// Set assignees for this file using the relative path if they were found
		path, err = filepath.Rel(o.projectDir, path)
		if err != nil {
			glog.Errorf("Unable to find relative path between %q and %q: %v", o.projectDir, path, err)
			return err
		}
		o.assignees[path] = sets.NewString()
		if len(c.Assignees) > 0 {
			o.assignees[path].Union(sets.NewString(c.Assignees...))
		}
		if len(c.Approvers) > 0 {
			o.assignees[path].Union(sets.NewString(c.Approvers...))
		}

		return nil
	}

	if filename != ownerFilename {
		return nil
	}

	file, err := os.Open(path)
	if err != nil {
		glog.Errorf("%v", err)
		return nil
	}
	defer file.Close()

	if err := yaml.NewYAMLToJSONDecoder(file).Decode(c); err != nil {
		glog.Errorf("%v", err)
		return nil
	}

	path, err = filepath.Rel(o.projectDir, path)
	if err != nil {
		glog.Errorf("Unable to find relative path between %q and %q: %v", o.projectDir, path, err)
		return err
	}
	path = filepath.Dir(path)
	// Make the root explicitly / so its easy to distinguish. Nothing else is `/` anchored
	if path == "." {
		path = "/"
	}
	o.assignees[path] = sets.NewString()
	if len(c.Assignees) > 0 {
		o.assignees[path].Union(sets.NewString(c.Assignees...))
	} else if len(c.Approvers) > 0 {
		o.assignees[path].Union(sets.NewString(c.Approvers...))
	}
	//if len(c.Owners) > 0 {
	//o.owners[path] = sets.NewString(c.Owners...)
	//}
	return nil
}

// decodeAssignmentConfig will parse the yaml header if it exists and unmarshal it into an assignmentConfig.
// If no yaml header is found, do nothing
// Returns an error if the file cannot be read or the yaml header is found but cannot be unmarshalled
var mdStructuredHeaderRegex = regexp.MustCompile("^---\n(.|\n)*\n---")

func decodeAssignmentConfig(path string, config *assignmentConfig) error {
	fileBytes, err := ioutil.ReadFile(path)
	if err != nil {
		return err
	}
	// Parse the yaml header from the top of the file.  Will return an empty string if regex does not match.
	meta := mdStructuredHeaderRegex.FindString(string(fileBytes))

	// Unmarshal the yaml header into the config
	return parseYaml.Unmarshal([]byte(meta), &config)
}

func (o *RepoInfo) updateRepoUsers() error {
	out, err := o.GitCommand([]string{"pull"})
	if err != nil {
		glog.Errorf("Unable to run git pull:\n%s\n%v", string(out), err)
		return err
	}

	out, err = o.GitCommand([]string{"rev-parse", "HEAD"})
	if err != nil {
		glog.Errorf("Unable get sha of HEAD:\n%s\n%v", string(out), err)
		return err
	}
	sha := out

	o.assignees = map[string]sets.String{}
	//o.owners = map[string]sets.String{}
	err = filepath.Walk(o.projectDir, o.walkFunc)
	if err != nil {
		glog.Errorf("Got error %v", err)
	}
	glog.Infof("Loaded config from %s:%s", o.projectDir, sha)
	glog.V(5).Infof("assignees: %v", o.assignees)
	//glog.V(5).Infof("owners: %v", o.owners)
	return nil
}

// Initialize will initialize the munger
func (o *RepoInfo) Initialize(config *github.Config) error {
	o.enabled = true
	o.config = config
	o.projectDir = path.Join(o.BaseDir, o.config.Project)

	if len(o.BaseDir) == 0 {
		glog.Fatalf("--repo-dir is required with selected munger(s)")
	}
	finfo, err := os.Stat(o.BaseDir)
	if err != nil {
		return fmt.Errorf("Unable to stat --repo-dir: %v", err)
	}
	if !finfo.IsDir() {
		return fmt.Errorf("--repo-dir is not a directory")
	}

	// check if the cloned dir already exists, if yes, cleanup.
	if _, err := os.Stat(o.projectDir); !os.IsNotExist(err) {
		if err := o.cleanUp(o.projectDir); err != nil {
			return fmt.Errorf("Unable to remove old clone directory at %v: %v", o.projectDir, err)
		}
	}

	if cloneUrl, err := o.cloneRepo(); err != nil {
		return fmt.Errorf("Unable to clone %v: %v", cloneUrl, err)
	}
	return o.updateRepoUsers()
}

func (o *RepoInfo) cleanUp(path string) error {
	return os.RemoveAll(path)
}

func (o *RepoInfo) cloneRepo() (string, error) {
	cloneUrl := fmt.Sprintf("https://github.com/%s/%s.git", o.config.Org, o.config.Project)
	output, err := o.gitCommandDir([]string{"clone", cloneUrl, o.projectDir}, o.BaseDir)
	if err != nil {
		glog.Errorf("Failed to clone github repo: %s", output)
	}
	return cloneUrl, err
}

// EachLoop is called at the start of every munge loop
func (o *RepoInfo) EachLoop() error {
	if !o.enabled {
		return nil
	}
	_, err := o.GitCommand([]string{"remote", "update"})
	if err != nil {
		glog.Errorf("Unable to git remote update: %v", err)
	}
	return o.updateRepoUsers()
}

// AddFlags will add any request flags to the cobra `cmd`
func (o *RepoInfo) AddFlags(cmd *cobra.Command) {
	cmd.Flags().StringVar(&o.BaseDir, "repo-dir", "", "Path to perform checkout of repository")
	cmd.Flags().BoolVar(&o.EnableMdYaml, "enable-md-yaml", false, "If true, look for assignees in md yaml headers.")
}

// GitCommand will execute the git command with the `args` within the project directory.
func (o *RepoInfo) GitCommand(args []string) ([]byte, error) {
	return o.gitCommandDir(args, o.projectDir)
}

// GitCommandDir will execute the git command with the `args` within the 'dir' directory.
func (o *RepoInfo) gitCommandDir(args []string, cmdDir string) ([]byte, error) {
	cmd := exec.Command("git", args...)
	cmd.Dir = cmdDir
	return cmd.CombinedOutput()
}

func peopleForPath(path string, people map[string]sets.String, leafOnly bool, enableMdYaml bool) sets.String {
	d := path
	if !enableMdYaml {
		d = filepath.Dir(path)
	}

	out := sets.NewString()
	for {
		// special case the root
		if d == "" {
			d = "/"
		}
		s, ok := people[d]
		if ok {
			out = out.Union(s)
			if leafOnly {
				break
			}
		}
		if d == "/" {
			break
		}
		d, _ = filepath.Split(d)
		d = strings.TrimSuffix(d, "/")
	}
	return out
}

// LeafAssignees returns a set of users who are the closest assignees to the
// requested file. If pkg/OWNERS has user1 and pkg/util/OWNERS has user2 this
// will only return user2 for the path pkg/util/sets/file.go
func (o *RepoInfo) LeafAssignees(path string) sets.String {
	return peopleForPath(path, o.assignees, true, o.EnableMdYaml)
}

// Assignees returns a set of users who are the closest assignees to the
// requested file. If pkg/OWNERS has user1 and pkg/util/OWNERS has user2 this
// will return both user1 and user2 for the path pkg/util/sets/file.go
func (o *RepoInfo) Assignees(path string) sets.String {
	return peopleForPath(path, o.assignees, false, o.EnableMdYaml)
}

//func (o *RepoInfo) LeafOwners(path string) sets.String {
//return people(path, o.owners, true)
//}

//func (o *RepoInfo) Owners(path string) sets.String {
//return people(path, o.owners, false)
//}
<|endoftext|>


<commit_before>package index

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/blevesearch/bleve"
	"github.com/docker/distribution"
	"github.com/docker/distribution/context"
	"github.com/docker/distribution/digest"
	ref "github.com/docker/distribution/reference"
	"github.com/docker/distribution/registry/client"
	"github.com/docker/docker/image"
	"github.com/docker/docker/reference"
	"github.com/docker/engine-api/types/container"
	"github.com/nhurel/dim/lib/registry"
	. "gopkg.in/check.v1"
	"io"
	"sync"
	"testing"
)

type RegistrySuite struct {
	index *Index
}

type NoOpRegistryClient struct {
	client.Registry
	//registry.RegistryClient
}

func (r *NoOpRegistryClient) Repositories(ctx context.Context, repos []string, last string) (int, error) {
	repos[0] = "httpd"
	repos[1] = "mysql"
	return 2, io.EOF
}

func (r *NoOpRegistryClient) NewRepository(parsedName reference.Named) (registry.Repository, error) {
	return &NoOpRegistryRepository{name: parsedName.Name()}, nil
}
func (r *NoOpRegistryClient) Search(query, advanced string) error {
	return nil
}

func (r *NoOpRegistryClient) WalkRepositories(repositories chan<- registry.Repository) error {
	return registry.WalkRepositories(r, repositories)
}

type NoOpRegistryRepository struct {
	distribution.Repository
	name string
}

func (r *NoOpRegistryRepository) AllTags() ([]string, error) {
	switch r.name {
	case "httpd":
		return []string{"2.2", "2.4"}, nil
	case "mysql":
		return []string{"5.5", "5.7"}, nil
	default:
		return []string{}, nil
	}
}
func (r *NoOpRegistryRepository) Image(tag string) (img *registry.Image, err error) {
	dg := fmt.Sprintf("%s:%s", r.name, tag)
	img = repoImages[dg]
	err = nil
	return
}
func (r *NoOpRegistryRepository) ImageFromManifest(tagDigest digest.Digest, digest string) (img *registry.Image, err error) {
	return nil, nil
}
func (r *NoOpRegistryRepository) DeleteImage(tag string) error {
	return nil
}

func (r *NoOpRegistryRepository) WalkImages(images chan<- *registry.Image) error {
	return registry.WalkImages(r, images)
}
func (r *NoOpRegistryRepository) Named() ref.Named {
	n, _ := ref.ParseNamed(r.name)
	return n
}

var repoImages = map[string]*registry.Image{
	"httpd:2.2": &registry.Image{
		&image.Image{
			V1Image: image.V1Image{
				Config: &container.Config{},
			},
		},
		"2.2",
		"httpd:2.2",
	},
	"httpd:2.4": &registry.Image{
		&image.Image{
			V1Image: image.V1Image{
				Config: &container.Config{},
			},
		},
		"2.4",
		"httpd:2.4",
	},
	"mysql:5.5": &registry.Image{
		&image.Image{
			V1Image: image.V1Image{
				Config: &container.Config{},
			},
		},
		"5.5",
		"mysql:5.5",
	},
	"mysql:5.7": &registry.Image{
		&image.Image{
			V1Image: image.V1Image{
				Config: &container.Config{},
			},
		},
		"5.7",
		"mysql:5.7",
	},
}

// Hook up gocheck into the "go test" runner.
func TestRegistry(t *testing.T) { TestingT(t) }

var _ = Suite(&RegistrySuite{})

func (s *RegistrySuite) SetUpSuite(c *C) {

	logrus.SetLevel(logrus.DebugLevel)
	fmt.Println("SetupSuite()")
	if i, err := MockIndex(); err != nil {
		logrus.WithError(err).Errorln("Failed to create index")
		return
	} else {
		fmt.Println("New index")
		s.index = &Index{i, "", nil, &NoOpRegistryClient{}, sync.WaitGroup{}}
	}

}

func (s *RegistrySuite) TearDownSuite(c *C) {
	//gock.Off()
}

func (s *RegistrySuite) TestBuild(c *C) {
	s.index.Build()
	fmt.Println("INDEX BUILT")
	srq := bleve.NewSearchRequest(bleve.NewMatchAllQuery())
	srs, err := s.index.Search(srq)
	c.Assert(err, IsNil)
	c.Assert(srs.Total, Equals, uint64(4))
}
<commit_msg>go vet<commit_after>package index

import (
	"fmt"
	"github.com/Sirupsen/logrus"
	"github.com/blevesearch/bleve"
	"github.com/docker/distribution"
	"github.com/docker/distribution/context"
	"github.com/docker/distribution/digest"
	ref "github.com/docker/distribution/reference"
	"github.com/docker/distribution/registry/client"
	"github.com/docker/docker/image"
	"github.com/docker/docker/reference"
	"github.com/docker/engine-api/types/container"
	"github.com/nhurel/dim/lib/registry"
	. "gopkg.in/check.v1"
	"io"
	"sync"
	"testing"
)

type RegistrySuite struct {
	index *Index
}

type NoOpRegistryClient struct {
	client.Registry
	//registry.RegistryClient
}

func (r *NoOpRegistryClient) Repositories(ctx context.Context, repos []string, last string) (int, error) {
	repos[0] = "httpd"
	repos[1] = "mysql"
	return 2, io.EOF
}

func (r *NoOpRegistryClient) NewRepository(parsedName reference.Named) (registry.Repository, error) {
	return &NoOpRegistryRepository{name: parsedName.Name()}, nil
}
func (r *NoOpRegistryClient) Search(query, advanced string) error {
	return nil
}

func (r *NoOpRegistryClient) WalkRepositories(repositories chan<- registry.Repository) error {
	return registry.WalkRepositories(r, repositories)
}

type NoOpRegistryRepository struct {
	distribution.Repository
	name string
}

func (r *NoOpRegistryRepository) AllTags() ([]string, error) {
	switch r.name {
	case "httpd":
		return []string{"2.2", "2.4"}, nil
	case "mysql":
		return []string{"5.5", "5.7"}, nil
	default:
		return []string{}, nil
	}
}
func (r *NoOpRegistryRepository) Image(tag string) (img *registry.Image, err error) {
	dg := fmt.Sprintf("%s:%s", r.name, tag)
	img = repoImages[dg]
	err = nil
	return
}
func (r *NoOpRegistryRepository) ImageFromManifest(tagDigest digest.Digest, digest string) (img *registry.Image, err error) {
	return nil, nil
}
func (r *NoOpRegistryRepository) DeleteImage(tag string) error {
	return nil
}

func (r *NoOpRegistryRepository) WalkImages(images chan<- *registry.Image) error {
	return registry.WalkImages(r, images)
}
func (r *NoOpRegistryRepository) Named() ref.Named {
	n, _ := ref.ParseNamed(r.name)
	return n
}

var repoImages = map[string]*registry.Image{
	"httpd:2.2": &registry.Image{
		Image: &image.Image{
			V1Image: image.V1Image{
				Config: &container.Config{},
			},
		},
		Tag:    "2.2",
		Digest: "httpd:2.2",
	},
	"httpd:2.4": &registry.Image{
		Image: &image.Image{
			V1Image: image.V1Image{
				Config: &container.Config{},
			},
		},
		Tag:    "2.4",
		Digest: "httpd:2.4",
	},
	"mysql:5.5": &registry.Image{
		Image: &image.Image{
			V1Image: image.V1Image{
				Config: &container.Config{},
			},
		},
		Tag:    "5.5",
		Digest: "mysql:5.5",
	},
	"mysql:5.7": &registry.Image{
		Image: &image.Image{
			V1Image: image.V1Image{
				Config: &container.Config{},
			},
		},
		Tag:    "5.7",
		Digest: "mysql:5.7",
	},
}

// Hook up gocheck into the "go test" runner.
func TestRegistry(t *testing.T) { TestingT(t) }

var _ = Suite(&RegistrySuite{})

func (s *RegistrySuite) SetUpSuite(c *C) {

	logrus.SetLevel(logrus.DebugLevel)
	fmt.Println("SetupSuite()")
	if i, err := MockIndex(); err != nil {
		logrus.WithError(err).Errorln("Failed to create index")
		return
	} else {
		fmt.Println("New index")
		s.index = &Index{i, "", nil, &NoOpRegistryClient{}, sync.WaitGroup{}}
	}

}

func (s *RegistrySuite) TearDownSuite(c *C) {
	//gock.Off()
}

func (s *RegistrySuite) TestBuild(c *C) {
	s.index.Build()
	fmt.Println("INDEX BUILT")
	srq := bleve.NewSearchRequest(bleve.NewMatchAllQuery())
	srs, err := s.index.Search(srq)
	c.Assert(err, IsNil)
	c.Assert(srs.Total, Equals, uint64(4))
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package gcsproxy_test

import (
	"testing"

	"github.com/jacobsa/gcloud/gcs/mock_gcs"
	"github.com/jacobsa/gcsfuse/gcsproxy"
	"github.com/jacobsa/gcsfuse/timeutil"
	. "github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
	"google.golang.org/cloud/storage"
)

func TestListingProxy(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Invariant-checking listing proxy
////////////////////////////////////////////////////////////////////////

// A wrapper around ListingProxy that calls CheckInvariants whenever invariants
// should hold. For catching logic errors early in the test.
type checkingListingProxy struct {
	wrapped *gcsproxy.ListingProxy
}

func (lp *checkingListingProxy) Name() string {
	lp.wrapped.CheckInvariants()
	defer lp.wrapped.CheckInvariants()
	return lp.wrapped.Name()
}

func (lp *checkingListingProxy) List() ([]*storage.Object, []string, error) {
	lp.wrapped.CheckInvariants()
	defer lp.wrapped.CheckInvariants()
	return lp.wrapped.List(context.Background())
}

func (lp *checkingListingProxy) NoteNewObject(o *storage.Object) error {
	lp.wrapped.CheckInvariants()
	defer lp.wrapped.CheckInvariants()
	return lp.wrapped.NoteNewObject(o)
}

func (lp *checkingListingProxy) NoteNewSubdirectory(name string) error {
	lp.wrapped.CheckInvariants()
	defer lp.wrapped.CheckInvariants()
	return lp.wrapped.NoteNewSubdirectory(name)
}

func (lp *checkingListingProxy) NoteRemoval(name string) error {
	lp.wrapped.CheckInvariants()
	defer lp.wrapped.CheckInvariants()
	return lp.wrapped.NoteRemoval(name)
}

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type ListingProxyTest struct {
	dirName string
	bucket  mock_gcs.MockBucket
	clock   timeutil.SimulatedClock
	lp      checkingListingProxy
}

var _ SetUpInterface = &ListingProxyTest{}

func init() { RegisterTestSuite(&ListingProxyTest{}) }

func (t *ListingProxyTest) SetUp(ti *TestInfo) {
	t.dirName = "some/dir/"
	t.bucket = mock_gcs.NewMockBucket(ti.MockController, "bucket")

	var err error
	t.lp.wrapped, err = gcsproxy.NewListingProxy(t.bucket, &t.clock, t.dirName)
	if err != nil {
		panic(err)
	}
}

////////////////////////////////////////////////////////////////////////
// Test functions
////////////////////////////////////////////////////////////////////////

func (t *ListingProxyTest) CreateForRootDirectory() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) CreateForIllegalDirectoryName() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) Name() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) List_CallsBucket() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) List_BucketFails() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) List_EmptyResult() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) List_NonEmptyResult() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) List_CacheHasExpired() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) NoteNewObject() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) NoteNewObject_IllegalNames() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) NoteNewSubdirectory() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) NoteNewSubdirectory_IllegalNames() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) NoteRemoval() {
	AssertTrue(false, "TODO")
}
<commit_msg>ListingProxyTest.CreateForIllegalDirectoryName<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package gcsproxy_test

import (
	"testing"

	"github.com/jacobsa/gcloud/gcs/mock_gcs"
	"github.com/jacobsa/gcsfuse/gcsproxy"
	"github.com/jacobsa/gcsfuse/timeutil"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/ogletest"
	"golang.org/x/net/context"
	"google.golang.org/cloud/storage"
)

func TestListingProxy(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Invariant-checking listing proxy
////////////////////////////////////////////////////////////////////////

// A wrapper around ListingProxy that calls CheckInvariants whenever invariants
// should hold. For catching logic errors early in the test.
type checkingListingProxy struct {
	wrapped *gcsproxy.ListingProxy
}

func (lp *checkingListingProxy) Name() string {
	lp.wrapped.CheckInvariants()
	defer lp.wrapped.CheckInvariants()
	return lp.wrapped.Name()
}

func (lp *checkingListingProxy) List() ([]*storage.Object, []string, error) {
	lp.wrapped.CheckInvariants()
	defer lp.wrapped.CheckInvariants()
	return lp.wrapped.List(context.Background())
}

func (lp *checkingListingProxy) NoteNewObject(o *storage.Object) error {
	lp.wrapped.CheckInvariants()
	defer lp.wrapped.CheckInvariants()
	return lp.wrapped.NoteNewObject(o)
}

func (lp *checkingListingProxy) NoteNewSubdirectory(name string) error {
	lp.wrapped.CheckInvariants()
	defer lp.wrapped.CheckInvariants()
	return lp.wrapped.NoteNewSubdirectory(name)
}

func (lp *checkingListingProxy) NoteRemoval(name string) error {
	lp.wrapped.CheckInvariants()
	defer lp.wrapped.CheckInvariants()
	return lp.wrapped.NoteRemoval(name)
}

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

type ListingProxyTest struct {
	dirName string
	bucket  mock_gcs.MockBucket
	clock   timeutil.SimulatedClock
	lp      checkingListingProxy
}

var _ SetUpInterface = &ListingProxyTest{}

func init() { RegisterTestSuite(&ListingProxyTest{}) }

func (t *ListingProxyTest) SetUp(ti *TestInfo) {
	t.dirName = "some/dir/"
	t.bucket = mock_gcs.NewMockBucket(ti.MockController, "bucket")

	var err error
	t.lp.wrapped, err = gcsproxy.NewListingProxy(t.bucket, &t.clock, t.dirName)
	if err != nil {
		panic(err)
	}
}

////////////////////////////////////////////////////////////////////////
// Test functions
////////////////////////////////////////////////////////////////////////

func (t *ListingProxyTest) CreateForRootDirectory() {
	_, err := gcsproxy.NewListingProxy(t.bucket, &t.clock, "")
	AssertEq(nil, err)
}

func (t *ListingProxyTest) CreateForIllegalDirectoryName() {
	_, err := gcsproxy.NewListingProxy(t.bucket, &t.clock, "foo/bar")

	AssertNe(nil, err)
	ExpectThat(err, Error(HasSubstr("foo/bar")))
	ExpectThat(err, Error(HasSubstr("directory name")))
}

func (t *ListingProxyTest) Name() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) List_CallsBucket() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) List_BucketFails() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) List_EmptyResult() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) List_NonEmptyResult() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) List_CacheHasExpired() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) NoteNewObject() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) NoteNewObject_IllegalNames() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) NoteNewSubdirectory() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) NoteNewSubdirectory_IllegalNames() {
	AssertTrue(false, "TODO")
}

func (t *ListingProxyTest) NoteRemoval() {
	AssertTrue(false, "TODO")
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"syscall"
	"time"

	"github.com/Code-Hex/golet"
)

func main() {
	cctx, cancel := context.WithCancel(context.Background())
	go func() {
		<-time.After(time.Second * 10)
		cancel()
	}()

	p := golet.New(cctx)
	p.EnableColor()

	// Execution interval
	p.SetInterval(time.Second * 1)

	// Send TERM signal when context cancel
	p.SetCtxCancelSignal(syscall.SIGTERM)

	p.Env(map[string]string{
		"NAME":  "codehex",
		"VALUE": "121",
	})

	p.Add(
		golet.Service{
			Exec: "plackup --port $PORT",
			Tag:  "plack",
		},
		golet.Service{
			Exec:   "echo 'This is cron!!'",
			Every:  "30 * * * * *",
			Worker: 2,
			Tag:    "cron",
		},
		golet.Service{
			Exec: `perl -E 'sleep 2; say "THIS IS DIE PROGRAM"; die;'`,
			Tag:  "DIE-Perl",
		},
	)

	p.Add(golet.Service{
		Code:   serveCode(),
		Worker: 3,
	})

	p.Run()
}

func serveCode() func(context.Context) error {
	return func(ctx context.Context) error {
		c := ctx.(*golet.Context)
		c.Println("Hello golet!! Port:", c.Port())
		mux := http.NewServeMux()
		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			fmt.Fprintf(w, "Hello, World")
			buf := strings.NewReader("This is log string\nNew line1\nNew line2\nNew line3\n")
			c.Copy(buf)
		})

		panicChan := make(chan interface{}, 1)
		srv := &http.Server{Addr: c.ServePort(), Handler: mux}

		// for ListenAndServe's error
		go func() {
			defer func() {
				if p := recover(); p != nil {
					panicChan <- p
				}
			}()
			if err := srv.ListenAndServe(); err != nil {
				panic(err)
			}
		}()

		// Wait channels
		for {
			select {
			// You can notify signal received.
			case <-c.Recv():
				signal, err := c.Signal()
				if err != nil {
					c.Println(err.Error())
					return err
				}
				switch signal {
				// If you send TERM or HUP, restart this callback immediately.
				case syscall.SIGTERM, syscall.SIGHUP:
					c.Println(signal.String())
					if err := srv.Shutdown(ctx); err != nil {
						return err
					}
					return fmt.Errorf("signal recieved SIGTERM, SIGHUP as error")
				// End of run
				case syscall.SIGINT:
					c.Println(signal.String())
					return nil
				}
			// To catch context cancel
			case <-ctx.Done():
				c.Println("context cancel")
				return srv.Shutdown(ctx)
			// panic for ListenAndServe
			case p := <-panicChan:
				return p.(error)
			}
		}
	}
}
<commit_msg>Updated example 2<commit_after>package main

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"syscall"
	"time"

	"github.com/Code-Hex/golet"
)

func main() {
	cctx, cancel := context.WithCancel(context.Background())
	go func() {
		<-time.After(time.Second * 10)
		cancel()
	}()

	p := golet.New(cctx)
	p.EnableColor()

	// Execution interval
	p.SetInterval(time.Second * 1)

	// Send TERM signal when context cancel
	p.SetCtxCancelSignal(syscall.SIGTERM)

	p.Env(map[string]string{
		"NAME":  "codehex",
		"VALUE": "121",
	})

	p.Add(
		golet.Service{
			Exec: "plackup --port $PORT",
			Tag:  "plack",
		},
		golet.Service{
			Exec:   "echo 'This is cron!!'",
			Every:  "30 * * * * *",
			Worker: 2,
			Tag:    "cron",
		},
		golet.Service{
			Exec: `perl -E 'sleep 2; say "THIS IS DIE PROGRAM"; die;'`,
			Tag:  "DIE-Perl",
		},
	)

	// You can comment out this
	p.Add(golet.Service{
		Code:   serveCode(),
		Worker: 3,
	})

	p.Run()
}

func serveCode() func(context.Context) error {
	return func(ctx context.Context) error {
		c := ctx.(*golet.Context)
		c.Println("Hello golet!! Port:", c.Port())
		mux := http.NewServeMux()
		mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
			fmt.Fprintf(w, "Hello, World")
			buf := strings.NewReader("This is log string\nNew line1\nNew line2\nNew line3\n")
			c.Copy(buf)
		})

		panicChan := make(chan interface{}, 1)
		srv := &http.Server{Addr: c.ServePort(), Handler: mux}

		// for ListenAndServe's error
		go func() {
			defer func() {
				if p := recover(); p != nil {
					panicChan <- p
				}
			}()
			if err := srv.ListenAndServe(); err != nil {
				panic(err)
			}
		}()

		// Wait channels
		for {
			select {
			// You can notify signal received
			case <-c.Recv():
				signal, err := c.Signal()
				if err != nil {
					c.Println(err.Error())
					return err
				}
				switch signal {
				// If you send TERM or HUP, restart this callback immediately
				case syscall.SIGTERM, syscall.SIGHUP:
					c.Println(signal.String())
					if err := srv.Shutdown(ctx); err != nil {
						return err
					}
					return fmt.Errorf("signal recieved SIGTERM, SIGHUP as error")
				// End of run
				case syscall.SIGINT:
					c.Println(signal.String())
					return nil
				}
			// To catch context cancel
			case <-ctx.Done():
				c.Println("context cancel")
				return srv.Shutdown(ctx)
			// panic for ListenAndServe
			case p := <-panicChan:
				return p.(error)
			}
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>Schedule next check before staring a check. If we schedule after start, we can easily get overruns.<commit_after><|endoftext|>


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

package sqlite3

import (
	"context"
	"database/sql"
	"fmt"
	"strings"

	"github.com/matrix-org/dendrite/internal/sqlutil"
	"github.com/matrix-org/dendrite/roomserver/storage/shared"
	"github.com/matrix-org/dendrite/roomserver/storage/tables"
	"github.com/matrix-org/dendrite/roomserver/types"
)

// TODO: previous_reference_sha256 was NOT NULL before but it broke sytest because
// sytest sends no SHA256 sums in the prev_events references in the soft-fail tests.
// In Postgres an empty BYTEA field is not NULL so it's fine there. In SQLite it
// seems to care that it's empty and therefore hits a NOT NULL constraint on insert.
// We should really work out what the right thing to do here is.
const previousEventSchema = `
  CREATE TABLE IF NOT EXISTS roomserver_previous_events (
    previous_event_id TEXT NOT NULL,
    previous_reference_sha256 BLOB,
    event_nids TEXT NOT NULL,
    UNIQUE (previous_event_id, previous_reference_sha256)
  );
`

// Insert an entry into the previous_events table.
// If there is already an entry indicating that an event references that previous event then
// add the event NID to the list to indicate that this event references that previous event as well.
// This should only be modified while holding a "FOR UPDATE" lock on the row in the rooms table for this room.
// The lock is necessary to avoid data races when checking whether an event is already referenced by another event.
const insertPreviousEventSQL = `
	INSERT OR REPLACE INTO roomserver_previous_events
	  (previous_event_id, previous_reference_sha256, event_nids)
	  VALUES ($1, $2, $3)
`

const selectPreviousEventNIDsSQL = `
	SELECT event_nids FROM roomserver_previous_events
	  WHERE previous_event_id = $1 AND previous_reference_sha256 = $2
`

// Check if the event is referenced by another event in the table.
// This should only be done while holding a "FOR UPDATE" lock on the row in the rooms table for this room.
const selectPreviousEventExistsSQL = `
	SELECT 1 FROM roomserver_previous_events
	  WHERE previous_event_id = $1 AND previous_reference_sha256 = $2
`

type previousEventStatements struct {
	db                            *sql.DB
	insertPreviousEventStmt       *sql.Stmt
	selectPreviousEventNIDsStmt   *sql.Stmt
	selectPreviousEventExistsStmt *sql.Stmt
}

func NewSqlitePrevEventsTable(db *sql.DB) (tables.PreviousEvents, error) {
	s := &previousEventStatements{
		db: db,
	}
	_, err := db.Exec(previousEventSchema)
	if err != nil {
		return nil, err
	}

	return s, shared.StatementList{
		{&s.insertPreviousEventStmt, insertPreviousEventSQL},
		{&s.selectPreviousEventNIDsStmt, selectPreviousEventNIDsSQL},
		{&s.selectPreviousEventExistsStmt, selectPreviousEventExistsSQL},
	}.Prepare(db)
}

func (s *previousEventStatements) InsertPreviousEvent(
	ctx context.Context,
	txn *sql.Tx,
	previousEventID string,
	previousEventReferenceSHA256 []byte,
	eventNID types.EventNID,
) error {
	var eventNIDs string
	eventNIDAsString := fmt.Sprintf("%d", eventNID)
	selectStmt := sqlutil.TxStmt(txn, s.selectPreviousEventExistsStmt)
	err := selectStmt.QueryRowContext(ctx, previousEventID, previousEventReferenceSHA256).Scan(&eventNIDs)
	if err != sql.ErrNoRows {
		return fmt.Errorf("selectStmt.QueryRowContext.Scan: %w", err)
	}
	var nids []string
	if eventNIDs != "" {
		nids = strings.Split(eventNIDs, ",")
		for _, nid := range nids {
			if nid == eventNIDAsString {
				return nil
			}
		}
		eventNIDs = strings.Join(append(nids, eventNIDAsString), ",")
	} else {
		eventNIDs = eventNIDAsString
	}
	insertStmt := sqlutil.TxStmt(txn, s.insertPreviousEventStmt)
	_, err = insertStmt.ExecContext(
		ctx, previousEventID, previousEventReferenceSHA256, eventNIDs,
	)
	return err
}

// Check if the event reference exists
// Returns sql.ErrNoRows if the event reference doesn't exist.
func (s *previousEventStatements) SelectPreviousEventExists(
	ctx context.Context, txn *sql.Tx, eventID string, eventReferenceSHA256 []byte,
) error {
	var ok int64
	stmt := sqlutil.TxStmt(txn, s.selectPreviousEventExistsStmt)
	return stmt.QueryRowContext(ctx, eventID, eventReferenceSHA256).Scan(&ok)
}
<commit_msg>Fix bug in error handling in SQLite InsertPreviousEvent (#1456)<commit_after>// Copyright 2017-2018 New Vector Ltd
// Copyright 2019-2020 The Matrix.org Foundation C.I.C.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package sqlite3

import (
	"context"
	"database/sql"
	"fmt"
	"strings"

	"github.com/matrix-org/dendrite/internal/sqlutil"
	"github.com/matrix-org/dendrite/roomserver/storage/shared"
	"github.com/matrix-org/dendrite/roomserver/storage/tables"
	"github.com/matrix-org/dendrite/roomserver/types"
)

// TODO: previous_reference_sha256 was NOT NULL before but it broke sytest because
// sytest sends no SHA256 sums in the prev_events references in the soft-fail tests.
// In Postgres an empty BYTEA field is not NULL so it's fine there. In SQLite it
// seems to care that it's empty and therefore hits a NOT NULL constraint on insert.
// We should really work out what the right thing to do here is.
const previousEventSchema = `
  CREATE TABLE IF NOT EXISTS roomserver_previous_events (
    previous_event_id TEXT NOT NULL,
    previous_reference_sha256 BLOB,
    event_nids TEXT NOT NULL,
    UNIQUE (previous_event_id, previous_reference_sha256)
  );
`

// Insert an entry into the previous_events table.
// If there is already an entry indicating that an event references that previous event then
// add the event NID to the list to indicate that this event references that previous event as well.
// This should only be modified while holding a "FOR UPDATE" lock on the row in the rooms table for this room.
// The lock is necessary to avoid data races when checking whether an event is already referenced by another event.
const insertPreviousEventSQL = `
	INSERT OR REPLACE INTO roomserver_previous_events
	  (previous_event_id, previous_reference_sha256, event_nids)
	  VALUES ($1, $2, $3)
`

const selectPreviousEventNIDsSQL = `
	SELECT event_nids FROM roomserver_previous_events
	  WHERE previous_event_id = $1 AND previous_reference_sha256 = $2
`

// Check if the event is referenced by another event in the table.
// This should only be done while holding a "FOR UPDATE" lock on the row in the rooms table for this room.
const selectPreviousEventExistsSQL = `
	SELECT 1 FROM roomserver_previous_events
	  WHERE previous_event_id = $1 AND previous_reference_sha256 = $2
`

type previousEventStatements struct {
	db                            *sql.DB
	insertPreviousEventStmt       *sql.Stmt
	selectPreviousEventNIDsStmt   *sql.Stmt
	selectPreviousEventExistsStmt *sql.Stmt
}

func NewSqlitePrevEventsTable(db *sql.DB) (tables.PreviousEvents, error) {
	s := &previousEventStatements{
		db: db,
	}
	_, err := db.Exec(previousEventSchema)
	if err != nil {
		return nil, err
	}

	return s, shared.StatementList{
		{&s.insertPreviousEventStmt, insertPreviousEventSQL},
		{&s.selectPreviousEventNIDsStmt, selectPreviousEventNIDsSQL},
		{&s.selectPreviousEventExistsStmt, selectPreviousEventExistsSQL},
	}.Prepare(db)
}

func (s *previousEventStatements) InsertPreviousEvent(
	ctx context.Context,
	txn *sql.Tx,
	previousEventID string,
	previousEventReferenceSHA256 []byte,
	eventNID types.EventNID,
) error {
	var eventNIDs string
	eventNIDAsString := fmt.Sprintf("%d", eventNID)
	selectStmt := sqlutil.TxStmt(txn, s.selectPreviousEventExistsStmt)
	err := selectStmt.QueryRowContext(ctx, previousEventID, previousEventReferenceSHA256).Scan(&eventNIDs)
	if err != nil && err != sql.ErrNoRows {
		return fmt.Errorf("selectStmt.QueryRowContext.Scan: %w", err)
	}
	var nids []string
	if eventNIDs != "" {
		nids = strings.Split(eventNIDs, ",")
		for _, nid := range nids {
			if nid == eventNIDAsString {
				return nil
			}
		}
		eventNIDs = strings.Join(append(nids, eventNIDAsString), ",")
	} else {
		eventNIDs = eventNIDAsString
	}
	insertStmt := sqlutil.TxStmt(txn, s.insertPreviousEventStmt)
	_, err = insertStmt.ExecContext(
		ctx, previousEventID, previousEventReferenceSHA256, eventNIDs,
	)
	return err
}

// Check if the event reference exists
// Returns sql.ErrNoRows if the event reference doesn't exist.
func (s *previousEventStatements) SelectPreviousEventExists(
	ctx context.Context, txn *sql.Tx, eventID string, eventReferenceSHA256 []byte,
) error {
	var ok int64
	stmt := sqlutil.TxStmt(txn, s.selectPreviousEventExistsStmt)
	return stmt.QueryRowContext(ctx, eventID, eventReferenceSHA256).Scan(&ok)
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"errors"
	"fmt"
	"sort"
)

const (
	evalInf = 50000
)

var (
	errCheckmate checkmateError
	errStalemate = errors.New("stalemate")
	errFiftyMove = errors.New("fifty-move rule")
	// TODO: threefold repetition
)

type Search struct {
	allowCutoff bool
	counters    []int
}

// SearchPosition searches a Position to the specified depth via iterative deepening
// and returns the search results.
func SearchPosition(ctx context.Context, pos Position, depth int) Results {
	var rs Results
	s := Search{
		allowCutoff: true,
		counters:    make([]int, depth+1),
	}
	for d := 1; d <= depth; d++ {
		_, rs, _ = s.negamax(ctx, pos, rs, Window{-evalInf, evalInf}, d)
		if ctx.Err() != nil {
			break
		}
	}
	return rs
}

// negamax recursively searches a Position to the specified depth and returns the evaluation score
// relative to the side to move and the search results. It employs alpha-beta pruning outside of
// the specified Window. If rs is zero length, negamax will generate and search all legal
// moves; if recommended moves are provided, they must all be legal, and only they will be searched.
func (s *Search) negamax(
	ctx context.Context,
	pos Position,
	rs Results,
	w Window,
	depth int,
) (bestScore RelScore, results Results, err error) {
	s.counters[len(s.counters)-1-depth]++

	score, rs, err := checkDone(pos, rs)
	if err != nil && (s.allowCutoff || err != errInsufficient) {
		// Do not cut off during perft in the case of insufficient material
		return score, rs, err
	}
	// Invariant: len(rs) > 0 and rs contains only legal moves
	if depth == 0 {
		score, err := Eval(pos)
		return score.Rel(pos.ToMove), rs, err
	}
	if s.allowCutoff && deepEnough(rs, depth) {
		return rs[0].score.Rel(pos.ToMove), rs, rs[0].err
	}

	defer func() { rs.SortFor(pos.ToMove) }()

	for _, r := range rs {
		if r.err != nil && s.allowCutoff {
			// The move is already known not to avoid a game-ending state; no need to search it further
			rs.Update(r)
			continue
		}

		score, cont, err := s.negamax(ctx, Make(pos, r.move), r.cont, w.Neg(), depth-1)
		score *= -1

		rs.Update(Result{move: r.move, score: score.Abs(pos.ToMove), depth: depth - 1, cont: cont, err: err})

		if depth >= 3 && ctx.Err() != nil {
			break
		}

		var constrained, ok bool
		w, constrained, ok = w.Constrain(score)
		if constrained {
			// improved lower bound
		}
		if !ok && s.allowCutoff {
			break
		}
	}
	if err, ok := rs[0].err.(checkmateError); ok {
		return w.alpha, rs, err.Prev()
	}
	return w.alpha, rs, rs[0].err
}

// checkDone reports whether pos represents a game-ending position.
// If so, the RelScore and error value indicate the type of ending.
// If not, the RelScore is 0 and the error is nil.
// In either case, checkDone returns a Results of all legal moves in pos.
// If recommended is not provided, checkDone generates the legal moves first.
func checkDone(pos Position, recommended Results) (RelScore, Results, error) {
	if len(recommended) == 0 {
		moves := PseudoLegalMoves(pos)
		recommended = make(Results, 0, len(moves))
		for _, m := range moves {
			if !IsLegal(Make(pos, m)) {
				continue
			}
			recommended = append(recommended, Result{move: m})
		}
	}
	if len(recommended) == 0 {
		// no legal moves
		if IsCheck(pos) {
			return -evalInf, rs, errCheckmate
		}
		return 0, rs, errStalemate
	}
	if pos.HalfMove == 100 {
		// fifty-move rule
		return 0, rs, errFiftyMove
	}
	return 0, recommended, nil
}

// deepEnough reports whether rs stores the results of a position search to at least the specified depth.
func deepEnough(rs Results, depth int) bool {
	// rs is already deep enough if all non-terminal elements have been searched to at least depth-1,
	// as long as they have been searched to non-zero depth
	for _, r := range rs {
		if (r.depth == 0 || r.depth < depth-1) && r.err == nil {
			return false
		}
	}
	return true
}

// IsPseudoLegal returns whether a Move is pseudo-legal in a Position.
// A move is pseudo-legal if the square to be moved from contains the specified piece
// and the piece is capable of moving to the target square if doing so would not put the king in check.
func IsPseudoLegal(pos Position, move Move) bool {
	for _, m := range PieceMoves[move.Piece](pos) {
		if m == move {
			return true
		}
	}
	return false
}

// IsLegal returns whether a Position is legal.
// A position is illegal if the king of the side that just moved is in check.
func IsLegal(pos Position) bool {
	return !IsAttacked(pos, pos.KingSquare[pos.Opp()], pos.ToMove)
}

// IsCheck returns whether the king of the side to move is in check.
func IsCheck(pos Position) bool {
	return IsAttacked(pos, pos.KingSquare[pos.ToMove], pos.Opp())
}

// IsTerminal returns whether or not a Position is checkmate or stalemate.
// A position is checkmate or stalemate if the side to move has no legal moves.
func IsTerminal(pos Position) bool { return !anyLegal(pos, PseudoLegalMoves(pos)) }

// anyLegal returns whether any of the given Moves are legal in the given Position.
func anyLegal(pos Position, moves []Move) bool {
	for _, m := range moves {
		if IsLegal(Make(pos, m)) {
			return true
		}
	}
	return false
}

// A Result holds a searched Move along with its evaluated Score,
// the search depth, and the continuation Results of the search.
type Result struct {
	move  Move
	score Score
	depth int
	cont  Results
	err   error
}

// String returns a string representation of r, including its principal variation.
func (r Result) String() string {
	s := fmt.Sprintf("%v (%v) %v", r.score, r.depth, r.PV())
	if r.err != nil {
		return s + " " + r.err.Error()
	}
	return s
}

// PV returns a string representation of r's principal variation.
func (r Result) PV() string {
	if r.depth == 0 || len(r.cont) == 0 {
		return LongAlgebraic(r.move)
	}
	return LongAlgebraic(r.move) + " " + r.cont[0].PV()
}

// A Results contains the results of a search.
// Functions returning a Results should sort it before returning.
type Results []Result

// SortFor sorts rs beginning with the best move for c.
// The sort is not guaranteed to be stable.
func (rs Results) SortFor(c Color) {
	// Sort first by mate condition, then by depth decreasing,
	// then by Score decreasing/increasing for White/Black,
	// and then by origin and destination Square increasing
	sort.Slice(rs, func(i, j int) bool {
		if less, ok := rs.mateSort(i, j); ok {
			return less
		}
		if rs[i].depth != rs[j].depth {
			return rs[i].depth > rs[j].depth
		}
		if rs[i].score != rs[j].score {
			return (rs[i].score > rs[j].score) == (c == White)
		}
		return rs.squareSort(i, j)
	})
}

// SortBySquares sorts rs first by the Result moves' origin Squares and then by their destination Squares,
// and finally by promotion piece in the case of pawn promotions.
func (rs Results) SortBySquares() { sort.Slice(rs, rs.squareSort) }

// mateSort reports how two Result elements should be sorted if at least one of them leads to checkmate.
// ok reports whether either Result contains an error of type checkmateError.
// If ok is true, less reports whether the Result with index i should sort before the Result with index j.
// If ok is false, the value of less is meaningless.
func (rs Results) mateSort(i, j int) (less bool, ok bool) {
	ich, iok := rs[i].err.(checkmateError)
	jch, jok := rs[j].err.(checkmateError)
	switch {
	case iok && jok:
		// Mates stored in rs as wins for the "previous player" are wins in the position being searched
		switch iwin, jwin := ich&1 == 0, jch&1 == 0; {
		case iwin && jwin:
			// faster winning mate first
			return ich < jch, true
		case !iwin && !jwin:
			// slower losing mate first
			return ich > jch, true
		default:
			return iwin, true
		}
	case iok:
		return ich&1 == 0, true
	case jok:
		return jch&1 != 0, true
	}
	return false, false
}

// squareSort reports how two Result elements should be sorted by their moves' origin and destination Squares,
// and then finally by promotion piece in the case of promoting pawns.
func (rs Results) squareSort(i, j int) bool {
	if rs[i].move.From != rs[j].move.From {
		return rs[i].move.From < rs[j].move.From
	}
	if rs[i].move.To != rs[j].move.To {
		return rs[i].move.To < rs[j].move.To
	}
	// From and To are non-unique in the case of pawn promotion
	return rs[i].move.PromotePiece < rs[j].move.PromotePiece
}

// Update finds the Result in rs with the same Move as r and replaces it with r.
// It panics if rs does not contain any elements with r's Move.
func (rs Results) Update(r Result) {
	for i := range rs {
		if rs[i].move == r.move {
			rs[i] = r
			return
		}
	}
	// rs should contain all legal moves
	panic("unreached")
}

// String returns a string representation of all Result values in r.
func (rs Results) String() string {
	var s string
	for _, r := range rs {
		s += fmt.Sprintf("%v\n", r.String())
	}
	return s
}

// A checkmateError indicates the number of plies until a forced checkmate can be delivered.
// Odd values are wins for the player with the next move, and even values for the player with the previous move.
// The zero value of type checkmateError indicates that the current position is checkmate.
type checkmateError int

func (n checkmateError) Error() string {
	if n&1 == 0 {
		return fmt.Sprintf("-#%d", n/2)
	}
	return fmt.Sprintf("#%d", (n+1)/2)
}

// Prev returns the checkmateError corresponding to the previous ply.
func (n checkmateError) Prev() checkmateError { return n + 1 }

// A Window represents the bounds of a position's evaluation.
type Window struct{ alpha, beta RelScore }

// Constrain updates the lower bound of w, if applicable, and returns the updated window,
// whether the lower bound was changed, and whether the returned Window remains valid.
func (w Window) Constrain(n RelScore) (c Window, constrained bool, ok bool) {
	if n <= w.alpha {
		return w, false, true
	}
	return Window{n, w.beta}, true, n <= w.beta
}

// Neg returns the additive inverse of w.
func (w Window) Neg() Window { return Window{-w.beta, -w.alpha} }
<commit_msg>correctly propagate checkmate errors<commit_after>package main

import (
	"context"
	"errors"
	"fmt"
	"sort"
)

const (
	evalInf = 50000
)

var (
	errCheckmate checkmateError
	errStalemate = errors.New("stalemate")
	errFiftyMove = errors.New("fifty-move rule")
	// TODO: threefold repetition
)

type Search struct {
	allowCutoff bool
	counters    []int
}

// SearchPosition searches a Position to the specified depth via iterative deepening
// and returns the search results.
func SearchPosition(ctx context.Context, pos Position, depth int) Results {
	var rs Results
	s := Search{
		allowCutoff: true,
		counters:    make([]int, depth+1),
	}
	for d := 1; d <= depth; d++ {
		_, rs, _ = s.negamax(ctx, pos, rs, Window{-evalInf, evalInf}, d)
		if ctx.Err() != nil {
			break
		}
	}
	return rs
}

// negamax recursively searches a Position to the specified depth and returns the evaluation score
// relative to the side to move and the search results. It employs alpha-beta pruning outside of
// the specified Window. If rs is zero length, negamax will generate and search all legal
// moves; if recommended moves are provided, they must all be legal, and only they will be searched.
func (s *Search) negamax(
	ctx context.Context,
	pos Position,
	rs Results,
	w Window,
	depth int,
) (bestScore RelScore, results Results, err error) {
	s.counters[len(s.counters)-1-depth]++

	score, rs, err := checkDone(pos, rs)
	if err != nil && (s.allowCutoff || err != errInsufficient) {
		// Do not cut off during perft in the case of insufficient material
		return score, rs, err
	}
	// Invariant: len(rs) > 0 and rs contains only legal moves
	if depth == 0 {
		score, err := Eval(pos)
		return score.Rel(pos.ToMove), rs, err
	}
	if s.allowCutoff && deepEnough(rs, depth) {
		if err, ok := rs[0].err.(checkmateError); ok {
			return rs[0].score.Rel(pos.ToMove), rs, err.Prev()
		}
		return rs[0].score.Rel(pos.ToMove), rs, rs[0].err
	}

	for _, r := range rs {
		if r.err != nil && s.allowCutoff {
			// The move is already known not to avoid a game-ending state; no need to search it further
			continue
		}

		score, cont, err := s.negamax(ctx, Make(pos, r.move), r.cont, w.Neg(), depth-1)
		score *= -1

		rs.Update(Result{move: r.move, score: score.Abs(pos.ToMove), depth: depth - 1, cont: cont, err: err})

		if depth >= 3 && ctx.Err() != nil {
			break
		}

		var constrained, ok bool
		w, constrained, ok = w.Constrain(score)
		if constrained {
			// improved lower bound
		}
		if !ok && s.allowCutoff {
			break
		}
	}
	rs.SortFor(pos.ToMove)
	return w.alpha, rs, rs[0].err
}

// checkDone reports whether pos represents a game-ending position.
// If so, the RelScore and error value indicate the type of ending.
// If not, the RelScore is 0 and the error is nil.
// In either case, checkDone returns a Results of all legal moves in pos.
// If rs is not provided, checkDone generates the legal moves first.
func checkDone(pos Position, rs Results) (RelScore, Results, error) {
	if len(rs) == 0 {
		moves := PseudoLegalMoves(pos)
		rs = make(Results, 0, len(moves))
		for _, m := range moves {
			if !IsLegal(Make(pos, m)) {
				continue
			}
			rs = append(rs, Result{move: m})
		}
	}
	if len(rs) == 0 {
		// no legal moves
		if IsCheck(pos) {
			return -evalInf, rs, errCheckmate
		}
		return 0, rs, errStalemate
	}
	if pos.HalfMove == 100 {
		// fifty-move rule
		return 0, rs, errFiftyMove
	}
	return 0, rs, nil
}

// deepEnough reports whether rs stores the results of a position search to at least the specified depth.
func deepEnough(rs Results, depth int) bool {
	// rs is already deep enough if all non-terminal elements have been searched to at least depth-1,
	// as long as they have been searched to non-zero depth
	for _, r := range rs {
		if (r.depth == 0 || r.depth < depth-1) && r.err == nil {
			return false
		}
	}
	return true
}

// IsPseudoLegal returns whether a Move is pseudo-legal in a Position.
// A move is pseudo-legal if the square to be moved from contains the specified piece
// and the piece is capable of moving to the target square if doing so would not put the king in check.
func IsPseudoLegal(pos Position, move Move) bool {
	for _, m := range PieceMoves[move.Piece](pos) {
		if m == move {
			return true
		}
	}
	return false
}

// IsLegal returns whether a Position is legal.
// A position is illegal if the king of the side that just moved is in check.
func IsLegal(pos Position) bool {
	return !IsAttacked(pos, pos.KingSquare[pos.Opp()], pos.ToMove)
}

// IsCheck returns whether the king of the side to move is in check.
func IsCheck(pos Position) bool {
	return IsAttacked(pos, pos.KingSquare[pos.ToMove], pos.Opp())
}

// IsTerminal returns whether or not a Position is checkmate or stalemate.
// A position is checkmate or stalemate if the side to move has no legal moves.
func IsTerminal(pos Position) bool { return !anyLegal(pos, PseudoLegalMoves(pos)) }

// anyLegal returns whether any of the given Moves are legal in the given Position.
func anyLegal(pos Position, moves []Move) bool {
	for _, m := range moves {
		if IsLegal(Make(pos, m)) {
			return true
		}
	}
	return false
}

// A Result holds a searched Move along with its evaluated Score,
// the search depth, and the continuation Results of the search.
type Result struct {
	move  Move
	score Score
	depth int
	cont  Results
	err   error
}

// String returns a string representation of r, including its principal variation.
func (r Result) String() string {
	s := fmt.Sprintf("%v (%v) %v", r.score, r.depth, r.PV())
	if r.err != nil {
		return s + " " + r.err.Error()
	}
	return s
}

// PV returns a string representation of r's principal variation.
func (r Result) PV() string {
	if r.depth == 0 || len(r.cont) == 0 {
		return LongAlgebraic(r.move)
	}
	return LongAlgebraic(r.move) + " " + r.cont[0].PV()
}

// A Results contains the results of a search.
// Functions returning a Results should sort it before returning.
type Results []Result

// SortFor sorts rs beginning with the best move for c.
// The sort is not guaranteed to be stable.
func (rs Results) SortFor(c Color) {
	// Sort first by mate condition, then by depth decreasing,
	// then by Score decreasing/increasing for White/Black,
	// and then by origin and destination Square increasing
	sort.Slice(rs, func(i, j int) bool {
		if less, ok := rs.mateSort(i, j); ok {
			return less
		}
		if rs[i].depth != rs[j].depth {
			return rs[i].depth > rs[j].depth
		}
		if rs[i].score != rs[j].score {
			return (rs[i].score > rs[j].score) == (c == White)
		}
		return rs.squareSort(i, j)
	})
}

// SortBySquares sorts rs first by the Result moves' origin Squares and then by their destination Squares,
// and finally by promotion piece in the case of pawn promotions.
func (rs Results) SortBySquares() { sort.Slice(rs, rs.squareSort) }

// mateSort reports how two Result elements should be sorted if at least one of them leads to checkmate.
// ok reports whether either Result contains an error of type checkmateError.
// If ok is true, less reports whether the Result with index i should sort before the Result with index j.
// If ok is false, the value of less is meaningless.
func (rs Results) mateSort(i, j int) (less bool, ok bool) {
	ich, iok := rs[i].err.(checkmateError)
	jch, jok := rs[j].err.(checkmateError)
	switch {
	case iok && jok:
		// Mates stored in rs as wins for the "previous player" are wins in the position being searched
		switch iwin, jwin := ich&1 == 0, jch&1 == 0; {
		case iwin && jwin:
			// faster winning mate first
			return ich < jch, true
		case !iwin && !jwin:
			// slower losing mate first
			return ich > jch, true
		default:
			return iwin, true
		}
	case iok:
		return ich&1 == 0, true
	case jok:
		return jch&1 != 0, true
	}
	return false, false
}

// squareSort reports how two Result elements should be sorted by their moves' origin and destination Squares,
// and then finally by promotion piece in the case of promoting pawns.
func (rs Results) squareSort(i, j int) bool {
	if rs[i].move.From != rs[j].move.From {
		return rs[i].move.From < rs[j].move.From
	}
	if rs[i].move.To != rs[j].move.To {
		return rs[i].move.To < rs[j].move.To
	}
	// From and To are non-unique in the case of pawn promotion
	return rs[i].move.PromotePiece < rs[j].move.PromotePiece
}

// Update finds the Result in rs with the same Move as r and replaces it with r.
// It panics if rs does not contain any elements with r's Move.
func (rs Results) Update(r Result) {
	for i := range rs {
		if rs[i].move == r.move {
			rs[i] = r
			return
		}
	}
	// rs should contain all legal moves
	panic("unreached")
}

// String returns a string representation of all Result values in r.
func (rs Results) String() string {
	var s string
	for _, r := range rs {
		s += fmt.Sprintf("%v\n", r.String())
	}
	return s
}

// A checkmateError indicates the number of plies until a forced checkmate can be delivered.
// Odd values are wins for the player with the next move, and even values for the player with the previous move.
// The zero value of type checkmateError indicates that the current position is checkmate.
type checkmateError int

func (n checkmateError) Error() string {
	if n&1 == 0 {
		return fmt.Sprintf("-#%d", n/2)
	}
	return fmt.Sprintf("#%d", (n+1)/2)
}

// Prev returns the checkmateError corresponding to the previous ply.
func (n checkmateError) Prev() checkmateError { return n + 1 }

// A Window represents the bounds of a position's evaluation.
type Window struct{ alpha, beta RelScore }

// Constrain updates the lower bound of w, if applicable, and returns the updated window,
// whether the lower bound was changed, and whether the returned Window remains valid.
func (w Window) Constrain(n RelScore) (c Window, constrained bool, ok bool) {
	if n <= w.alpha {
		return w, false, true
	}
	return Window{n, w.beta}, true, n <= w.beta
}

// Neg returns the additive inverse of w.
func (w Window) Neg() Window { return Window{-w.beta, -w.alpha} }
<|endoftext|>