<commit_before>package irmaserver

import (
	//TODO: use redigo instead of redis-go v8?
	"context"
	"encoding/json"
	"github.com/go-errors/errors"
	"strings"
	"sync"
	"time"

	"github.com/alexandrevicenzi/go-sse"
	"github.com/go-redis/redis/v8"
	"github.com/privacybydesign/gabi"
	"github.com/privacybydesign/gabi/big"
	irma "github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/internal/common"
	"github.com/privacybydesign/irmago/server"

	"github.com/sirupsen/logrus"
)

type session struct {
	//TODO: check if we can get rid of this Mutex for Redis
	sync.Mutex `json:-`
	//TODO: note somewhere that state with redis will not support sse for the moment
	sse         *sse.Server
	locked      bool
	sessions    sessionStore
	conf        *server.Configuration
	request     irma.SessionRequest
	context     context.Context
	toBeUpdated bool

	sessionData
}

type sessionData struct {
	Action             irma.Action
	Token              string
	ClientToken        string
	Version            *irma.ProtocolVersion `json:",omitempty"`
	Rrequest           irma.RequestorRequest
	LegacyCompatible   bool // if the request is convertible to pre-condiscon format
	ImplicitDisclosure irma.AttributeConDisCon
	Status             server.Status
	PrevStatus         server.Status
	ResponseCache      responseCache
	LastActive         time.Time
	Result             *server.SessionResult
	KssProofs          map[irma.SchemeManagerIdentifier]*gabi.ProofP
}

type responseCache struct {
	Message       []byte
	Response      []byte
	Status        int
	SessionStatus server.Status
}

type sessionStore interface {
	get(token string) (*session, error)
	clientGet(token string) (*session, error)
	add(session *session) error
	update(session *session) error
	stop()
}

type memorySessionStore struct {
	sync.RWMutex
	conf *server.Configuration

	requestor map[string]*session
	client    map[string]*session
}

type redisSessionStore struct {
	client *redis.Client
	conf   *server.Configuration
}

type RedisError interface {
	Error() string
}

const (
	maxSessionLifetime         = 5 * time.Minute // After this a session is cancelled
	requestorTokenLookupPrefix = "token:"
	clientTokenLookupPrefix    = "session:"
)

var (
	minProtocolVersion = irma.NewVersion(2, 4)
	maxProtocolVersion = irma.NewVersion(2, 7)
)

func (s *memorySessionStore) get(t string) (*session, error) {
	s.RLock()
	defer s.RUnlock()
	return s.requestor[t], nil
}

func (s *memorySessionStore) clientGet(t string) (*session, error) {
	s.RLock()
	defer s.RUnlock()
	return s.client[t], nil
}

func (s *memorySessionStore) add(session *session) error {
	s.Lock()
	defer s.Unlock()
	s.requestor[session.Token] = session
	s.client[session.ClientToken] = session
	return nil
}

func (s *memorySessionStore) update(_ *session) error {
	return nil
}

func (s *memorySessionStore) stop() {
	s.Lock()
	defer s.Unlock()
	for _, session := range s.requestor {
		if session.sse != nil {
			session.sse.CloseChannel("session/" + session.Token)
			session.sse.CloseChannel("session/" + session.ClientToken)
		}
	}
}

func (s *memorySessionStore) deleteExpired() {
	// First check which sessions have expired
	// We don't need a write lock for this yet, so postpone that for actual deleting
	s.RLock()
	expired := make([]string, 0, len(s.requestor))
	for token, session := range s.requestor {
		session.Lock()

		timeout := maxSessionLifetime
		if session.Status == server.StatusInitialized && session.Rrequest.Base().ClientTimeout != 0 {
			timeout = time.Duration(session.Rrequest.Base().ClientTimeout) * time.Second
		}

		if session.LastActive.Add(timeout).Before(time.Now()) {
			if !session.Status.Finished() {
				s.conf.Logger.WithFields(logrus.Fields{"session": session.Token}).Infof("Session expired")
				session.markAlive()
				session.setStatus(server.StatusTimeout)
			} else {
				s.conf.Logger.WithFields(logrus.Fields{"session": session.Token}).Infof("Deleting session")
				expired = append(expired, token)
			}
		}
		session.Unlock()
	}
	s.RUnlock()

	// Using a write lock, delete the expired sessions
	s.Lock()
	for _, token := range expired {
		session := s.requestor[token]
		if session.sse != nil {
			session.sse.CloseChannel("session/" + session.Token)
			session.sse.CloseChannel("session/" + session.ClientToken)
		}
		delete(s.client, session.ClientToken)
		delete(s.requestor, token)
	}
	s.Unlock()
}

// UnmarshalJSON unmarshals sessionData.
func (s *sessionData) UnmarshalJSON(data []byte) error {
	type rawSessionData sessionData

	var temp struct {
		Rrequest *json.RawMessage `json:",omitempty"`
		rawSessionData
	}

	if err := json.Unmarshal(data, &temp); err != nil {
		return err
	}

	*s = sessionData(temp.rawSessionData)

	if temp.Rrequest == nil {
		s.Rrequest = nil
		return errors.Errorf("temp.Rrequest == nil: %d \n", temp.Rrequest)
	}

	// unmarshal Rrequest
	ipR := &irma.IdentityProviderRequest{}
	spR := &irma.ServiceProviderRequest{}
	sigR := &irma.SignatureRequestorRequest{}

	if err := json.Unmarshal(*temp.Rrequest, ipR); err == nil && s.Action == "issuing" {
		s.Rrequest = ipR
	} else if err = json.Unmarshal(*temp.Rrequest, spR); err == nil && s.Action == "disclosing" {
		s.Rrequest = spR
	} else if err = json.Unmarshal(*temp.Rrequest, sigR); err == nil && s.Action == "signing" {
		s.Rrequest = sigR
	} else {
		return err
	}

	return nil
}

func (s *redisSessionStore) get(t string) (*session, error) {
	//TODO: input validation string?
	val, err := s.client.Get(context.Background(), requestorTokenLookupPrefix+t).Result()
	if err == redis.Nil {
		return nil, nil
	} else if err != nil {
		return nil, logAsRedisError(err)
	}

	return s.clientGet(val)
}

func (s *redisSessionStore) clientGet(t string) (*session, error) {
	val, err := s.client.Get(context.Background(), clientTokenLookupPrefix+t).Result()
	if err == redis.Nil {
		return nil, nil
	} else if err != nil {
		return nil, logAsRedisError(err)
	}

	var session session
	session.conf = s.conf
	session.sessions = s
	if err := json.Unmarshal([]byte(val), &session.sessionData); err != nil {
		return nil, logAsRedisError(err)
	}
	session.request = session.Rrequest.SessionRequest()

	if session.LastActive.Add(maxSessionLifetime).Before(time.Now()) && !session.Status.Finished() {
		s.conf.Logger.WithFields(logrus.Fields{"session": session.Token}).Infof("Session expired")
		session.markAlive()
		session.setStatus(server.StatusTimeout)
		_ = s.update(&session) // Worst case the TTL and status aren't updated. We won't deal with this error
	}

	return &session, nil
}

func (s *redisSessionStore) add(session *session) error {
	timeout := 2 * maxSessionLifetime // logic similar to memory store
	if session.Status == server.StatusInitialized && session.Rrequest.Base().ClientTimeout != 0 {
		timeout = time.Duration(session.Rrequest.Base().ClientTimeout) * time.Second
	} else if session.Status.Finished() {
		timeout = maxSessionLifetime
	}

	sessionJSON, err := json.Marshal(session.sessionData)
	if err != nil {
		return logAsRedisError(err)
	}

	err = s.client.Set(session.context, requestorTokenLookupPrefix+session.sessionData.Token, session.sessionData.ClientToken, timeout).Err()
	if err != nil {
		return logAsRedisError(err)
	}
	err = s.client.Set(session.context, clientTokenLookupPrefix+session.sessionData.ClientToken, sessionJSON, timeout).Err()
	if err != nil {
		return logAsRedisError(err)
	}

	return nil
}

func (s *redisSessionStore) update(session *session) error {
	return s.add(session)
}

func (s *redisSessionStore) stop() {
	err := s.client.Close()
	if err != nil {
		_ = logAsRedisError(err)
	}
}

var one *big.Int = big.NewInt(1)

func (s *Server) newSession(action irma.Action, request irma.RequestorRequest, ctx context.Context) (*session, error) {
	token := common.NewSessionToken()
	clientToken := common.NewSessionToken()

	base := request.SessionRequest().Base()
	if s.conf.AugmentClientReturnURL && base.AugmentReturnURL && base.ClientReturnURL != "" {
		if strings.Contains(base.ClientReturnURL, "?") {
			base.ClientReturnURL += "&token=" + token
		} else {
			base.ClientReturnURL += "?token=" + token
		}
	}

	sd := sessionData{
		Action:      action,
		Rrequest:    request,
		LastActive:  time.Now(),
		Token:       token,
		ClientToken: clientToken,
		Status:      server.StatusInitialized,
		PrevStatus:  server.StatusInitialized,
		Result: &server.SessionResult{
			LegacySession: request.SessionRequest().Base().Legacy(),
			Token:         token,
			Type:          action,
			Status:        server.StatusInitialized,
		},
	}
	ses := &session{
		sessionData: sd,
		sessions:    s.sessions,
		sse:         s.serverSentEvents,
		conf:        s.conf,
		request:     request.SessionRequest(),
		context:     ctx,
	}

	s.conf.Logger.WithFields(logrus.Fields{"session": ses.Token}).Debug("New session started")
	nonce, _ := gabi.GenerateNonce()
	base.Nonce = nonce
	base.Context = one
	err := s.sessions.add(ses)
	if err != nil {
		return nil, err
	}

	return ses, nil
}

func logAsRedisError(err error) error {
	return server.LogError(RedisError(err))
}
<commit_msg>refactor: improve code quality when unmarshaling sessionData.<commit_after>package irmaserver

import (
	//TODO: use redigo instead of redis-go v8?
	"context"
	"encoding/json"
	"github.com/go-errors/errors"
	"strings"
	"sync"
	"time"

	"github.com/alexandrevicenzi/go-sse"
	"github.com/go-redis/redis/v8"
	"github.com/privacybydesign/gabi"
	"github.com/privacybydesign/gabi/big"
	irma "github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/internal/common"
	"github.com/privacybydesign/irmago/server"

	"github.com/sirupsen/logrus"
)

type session struct {
	//TODO: check if we can get rid of this Mutex for Redis
	sync.Mutex `json:-`
	//TODO: note somewhere that state with redis will not support sse for the moment
	sse         *sse.Server
	locked      bool
	sessions    sessionStore
	conf        *server.Configuration
	request     irma.SessionRequest
	context     context.Context
	toBeUpdated bool

	sessionData
}

type sessionData struct {
	Action             irma.Action
	Token              string
	ClientToken        string
	Version            *irma.ProtocolVersion `json:",omitempty"`
	Rrequest           irma.RequestorRequest
	LegacyCompatible   bool // if the request is convertible to pre-condiscon format
	ImplicitDisclosure irma.AttributeConDisCon
	Status             server.Status
	PrevStatus         server.Status
	ResponseCache      responseCache
	LastActive         time.Time
	Result             *server.SessionResult
	KssProofs          map[irma.SchemeManagerIdentifier]*gabi.ProofP
}

type responseCache struct {
	Message       []byte
	Response      []byte
	Status        int
	SessionStatus server.Status
}

type sessionStore interface {
	get(token string) (*session, error)
	clientGet(token string) (*session, error)
	add(session *session) error
	update(session *session) error
	stop()
}

type memorySessionStore struct {
	sync.RWMutex
	conf *server.Configuration

	requestor map[string]*session
	client    map[string]*session
}

type redisSessionStore struct {
	client *redis.Client
	conf   *server.Configuration
}

type RedisError interface {
	Error() string
}

const (
	maxSessionLifetime         = 5 * time.Minute // After this a session is cancelled
	requestorTokenLookupPrefix = "token:"
	clientTokenLookupPrefix    = "session:"
)

var (
	minProtocolVersion = irma.NewVersion(2, 4)
	maxProtocolVersion = irma.NewVersion(2, 7)
)

func (s *memorySessionStore) get(t string) (*session, error) {
	s.RLock()
	defer s.RUnlock()
	return s.requestor[t], nil
}

func (s *memorySessionStore) clientGet(t string) (*session, error) {
	s.RLock()
	defer s.RUnlock()
	return s.client[t], nil
}

func (s *memorySessionStore) add(session *session) error {
	s.Lock()
	defer s.Unlock()
	s.requestor[session.Token] = session
	s.client[session.ClientToken] = session
	return nil
}

func (s *memorySessionStore) update(_ *session) error {
	return nil
}

func (s *memorySessionStore) stop() {
	s.Lock()
	defer s.Unlock()
	for _, session := range s.requestor {
		if session.sse != nil {
			session.sse.CloseChannel("session/" + session.Token)
			session.sse.CloseChannel("session/" + session.ClientToken)
		}
	}
}

func (s *memorySessionStore) deleteExpired() {
	// First check which sessions have expired
	// We don't need a write lock for this yet, so postpone that for actual deleting
	s.RLock()
	expired := make([]string, 0, len(s.requestor))
	for token, session := range s.requestor {
		session.Lock()

		timeout := maxSessionLifetime
		if session.Status == server.StatusInitialized && session.Rrequest.Base().ClientTimeout != 0 {
			timeout = time.Duration(session.Rrequest.Base().ClientTimeout) * time.Second
		}

		if session.LastActive.Add(timeout).Before(time.Now()) {
			if !session.Status.Finished() {
				s.conf.Logger.WithFields(logrus.Fields{"session": session.Token}).Infof("Session expired")
				session.markAlive()
				session.setStatus(server.StatusTimeout)
			} else {
				s.conf.Logger.WithFields(logrus.Fields{"session": session.Token}).Infof("Deleting session")
				expired = append(expired, token)
			}
		}
		session.Unlock()
	}
	s.RUnlock()

	// Using a write lock, delete the expired sessions
	s.Lock()
	for _, token := range expired {
		session := s.requestor[token]
		if session.sse != nil {
			session.sse.CloseChannel("session/" + session.Token)
			session.sse.CloseChannel("session/" + session.ClientToken)
		}
		delete(s.client, session.ClientToken)
		delete(s.requestor, token)
	}
	s.Unlock()
}

// UnmarshalJSON unmarshals sessionData.
func (s *sessionData) UnmarshalJSON(data []byte) error {
	type rawSessionData sessionData

	var temp struct {
		Rrequest *json.RawMessage `json:",omitempty"`
		rawSessionData
	}

	if err := json.Unmarshal(data, &temp); err != nil {
		return err
	}

	*s = sessionData(temp.rawSessionData)

	if temp.Rrequest == nil {
		s.Rrequest = nil
		return errors.Errorf("temp.Rrequest == nil: %d \n", temp.Rrequest)
	}

	// unmarshal Rrequest
	switch s.Action {
	case "issuing":
		s.Rrequest = &irma.IdentityProviderRequest{}
	case "disclosing":
		s.Rrequest = &irma.ServiceProviderRequest{}
	case "signing":
		s.Rrequest = &irma.SignatureRequestorRequest{}
	}

	return json.Unmarshal(*temp.Rrequest, s.Rrequest)
}

func (s *redisSessionStore) get(t string) (*session, error) {
	//TODO: input validation string?
	val, err := s.client.Get(context.Background(), requestorTokenLookupPrefix+t).Result()
	if err == redis.Nil {
		return nil, nil
	} else if err != nil {
		return nil, logAsRedisError(err)
	}

	return s.clientGet(val)
}

func (s *redisSessionStore) clientGet(t string) (*session, error) {
	val, err := s.client.Get(context.Background(), clientTokenLookupPrefix+t).Result()
	if err == redis.Nil {
		return nil, nil
	} else if err != nil {
		return nil, logAsRedisError(err)
	}

	var session session
	session.conf = s.conf
	session.sessions = s
	if err := json.Unmarshal([]byte(val), &session.sessionData); err != nil {
		return nil, logAsRedisError(err)
	}
	session.request = session.Rrequest.SessionRequest()

	if session.LastActive.Add(maxSessionLifetime).Before(time.Now()) && !session.Status.Finished() {
		s.conf.Logger.WithFields(logrus.Fields{"session": session.Token}).Infof("Session expired")
		session.markAlive()
		session.setStatus(server.StatusTimeout)
		_ = s.update(&session) // Worst case the TTL and status aren't updated. We won't deal with this error
	}

	return &session, nil
}

func (s *redisSessionStore) add(session *session) error {
	timeout := 2 * maxSessionLifetime // logic similar to memory store
	if session.Status == server.StatusInitialized && session.Rrequest.Base().ClientTimeout != 0 {
		timeout = time.Duration(session.Rrequest.Base().ClientTimeout) * time.Second
	} else if session.Status.Finished() {
		timeout = maxSessionLifetime
	}

	sessionJSON, err := json.Marshal(session.sessionData)
	if err != nil {
		return logAsRedisError(err)
	}

	err = s.client.Set(session.context, requestorTokenLookupPrefix+session.sessionData.Token, session.sessionData.ClientToken, timeout).Err()
	if err != nil {
		return logAsRedisError(err)
	}
	err = s.client.Set(session.context, clientTokenLookupPrefix+session.sessionData.ClientToken, sessionJSON, timeout).Err()
	if err != nil {
		return logAsRedisError(err)
	}

	return nil
}

func (s *redisSessionStore) update(session *session) error {
	return s.add(session)
}

func (s *redisSessionStore) stop() {
	err := s.client.Close()
	if err != nil {
		_ = logAsRedisError(err)
	}
}

var one *big.Int = big.NewInt(1)

func (s *Server) newSession(action irma.Action, request irma.RequestorRequest, ctx context.Context) (*session, error) {
	token := common.NewSessionToken()
	clientToken := common.NewSessionToken()

	base := request.SessionRequest().Base()
	if s.conf.AugmentClientReturnURL && base.AugmentReturnURL && base.ClientReturnURL != "" {
		if strings.Contains(base.ClientReturnURL, "?") {
			base.ClientReturnURL += "&token=" + token
		} else {
			base.ClientReturnURL += "?token=" + token
		}
	}

	sd := sessionData{
		Action:      action,
		Rrequest:    request,
		LastActive:  time.Now(),
		Token:       token,
		ClientToken: clientToken,
		Status:      server.StatusInitialized,
		PrevStatus:  server.StatusInitialized,
		Result: &server.SessionResult{
			LegacySession: request.SessionRequest().Base().Legacy(),
			Token:         token,
			Type:          action,
			Status:        server.StatusInitialized,
		},
	}
	ses := &session{
		sessionData: sd,
		sessions:    s.sessions,
		sse:         s.serverSentEvents,
		conf:        s.conf,
		request:     request.SessionRequest(),
		context:     ctx,
	}

	s.conf.Logger.WithFields(logrus.Fields{"session": ses.Token}).Debug("New session started")
	nonce, _ := gabi.GenerateNonce()
	base.Nonce = nonce
	base.Context = one
	err := s.sessions.add(ses)
	if err != nil {
		return nil, err
	}

	return ses, nil
}

func logAsRedisError(err error) error {
	return server.LogError(RedisError(err))
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"text/template"
)

type XapiTypes struct {
	XapiObjects []XapiType
}

type XapiType struct {
	Name        XapiName
	Description string
	Fields      []XapiField
	Messages    []XapiMessage
}

type XapiField struct {
	Name        XapiName
	Description string
	Type        XapiTypeName
	Qualifier   string
}

type XapiMessage struct {
	Name        string
	Description string
	Params      []XapiParam
	Result      []string
}

type XapiParam struct {
	Name string
	Type XapiTypeName
	Doc  string
}

func (message XapiMessage) ParamsLen() int {
	return len(message.Params)
}
func (name XapiName) String() string {
	s := string(name)

	s = strings.Replace(s, "_h", "H", -1)
	s = strings.Replace(s, "_u", "U", -1)
	s = strings.Replace(s, "_a", "A", -1)
	s = strings.Replace(s, "_c", "C", -1)
	s = strings.Replace(s, "_l", "L", -1)
	s = strings.Replace(s, "_s", "S", -1)
	s = strings.Replace(s, "_n", "N", -1)
	s = strings.Replace(s, "_t", "T", -1)
	s = strings.Replace(s, "_p", "P", -1)

	return strings.Title(s)
}

func (name XapiTypeName) String() string {
	s := string(name)
	if s == "string" || s == "bool" {
		return s
	}
	s = strings.Replace(s, " ", "_", -1)

	if s == "(string_->_string)_map" {
		return "map[string]string"
	}

	return "string // " + s
}

type XapiName string
type XapiTypeName string

func main() {
	file, e := ioutil.ReadFile("./xapi.json")
	if e != nil {
		fmt.Printf("File error: %v\n", e)
		os.Exit(1)
	}
	var types XapiTypes
	json.Unmarshal(file, &types)

	t, err := template.New("xapi.template").ParseFiles("xapi.template")
	if err != nil {
		fmt.Printf("execution failed: %s\n", err)
	}

	err = t.ExecuteTemplate(os.Stdout, "xapi.template", types)
	if err != nil {
		fmt.Printf("execution failed: %s\n", err)
	}
}
<commit_msg>renamings<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
	"text/template"
)

type Xapi struct {
	XapiObjects []XapiType
}

type XapiType struct {
	Name        XapiName
	Description string
	Fields      []XapiField
	Messages    []XapiMessage
}

type XapiField struct {
	Name        XapiName
	Description string
	Type        XapiTypeName
	Qualifier   string
}

type XapiMessage struct {
	Name        string
	Description string
	Params      []XapiParam
	Result      []string
}

type XapiParam struct {
	Name string
	Type XapiTypeName
	Doc  string
}

func (message XapiMessage) ParamsLen() int {
	return len(message.Params)
}
func (name XapiName) String() string {
	s := string(name)

	s = strings.Replace(s, "_h", "H", -1)
	s = strings.Replace(s, "_u", "U", -1)
	s = strings.Replace(s, "_a", "A", -1)
	s = strings.Replace(s, "_c", "C", -1)
	s = strings.Replace(s, "_l", "L", -1)
	s = strings.Replace(s, "_s", "S", -1)
	s = strings.Replace(s, "_n", "N", -1)
	s = strings.Replace(s, "_t", "T", -1)
	s = strings.Replace(s, "_p", "P", -1)

	return strings.Title(s)
}

func (name XapiTypeName) String() string {
	s := string(name)
	if s == "string" || s == "bool" {
		return s
	}
	s = strings.Replace(s, " ", "_", -1)

	if s == "(string_->_string)_map" {
		return "map[string]string"
	}

	return "string // " + s
}

type XapiName string
type XapiTypeName string

func main() {
	file, e := ioutil.ReadFile("./xapi.json")
	if e != nil {
		fmt.Printf("File error: %v\n", e)
		os.Exit(1)
	}
	var xapi Xapi
	json.Unmarshal(file, &xapi)

	t, err := template.New("xapi.template").ParseFiles("xapi.template")
	if err != nil {
		fmt.Printf("execution failed: %s\n", err)
	}

	err = t.ExecuteTemplate(os.Stdout, "xapi.template", xapi)
	if err != nil {
		fmt.Printf("execution failed: %s\n", err)
	}
}
<|endoftext|>


<commit_before>package dragoon

import (
	"time"

	"github.com/pkg/errors"
	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
)

type (
	// Identifier gives getter and setter for ID.
	Identifier interface {
		GetID() string
		SetID(string)
	}
	// TimeStamper gives getter for CreatedAt, setter for CreatedAt and UpdatedAt.
	TimeStamper interface {
		GetCreatedAt() time.Time
		SetCreatedAt(time.Time)
		SetUpdatedAt(time.Time)
	}
	// IdentifyGenerator gives generate of ID method.
	IdentifyGenerator interface {
		NextID(c context.Context, kind string) (string, error)
	}
	// Validator gives validate of fields method.
	Validator interface {
		Struct(target interface{}) error
	}
	// Spear has convenience methods.
	Spear struct {
		kind                string
		ignoreFieldMismatch bool
		identifyGenerator   IdentifyGenerator
		validator           Validator
	}
)

// ErrConflictEntity is returned when an entity was conflict for a given key.
var ErrConflictEntity = errors.New("dragoon: conflict entity")

// NewSpear returns new Spear.
func NewSpear(kind string, ignoreFieldMismatch bool, i IdentifyGenerator, v Validator) (*Spear, error) {
	if kind == "" || i == nil || v == nil {
		return nil, errors.New("dragoon: invalid arguments")
	}
	return &Spear{
		kind:                kind,
		ignoreFieldMismatch: ignoreFieldMismatch,
		identifyGenerator:   i,
		validator:           v,
	}, nil
}

// Get loads the entity based on e's key into e.
func (s *Spear) Get(c context.Context, e Identifier) error {
	k := datastore.NewKey(c, s.kind, e.GetID(), 0, nil)
	err := datastore.Get(c, k, e)
	if err != nil {
		if s.ignoreFieldMismatch && IsErrFieldMismatch(err) {
			return nil
		}
		return errors.Wrapf(err, "dragoon: failed to get an entity - key = %#v", k)
	}
	return nil
}

// GetMulti is a batch version of Get.
func (s *Spear) GetMulti(c context.Context, es []Identifier) error {
	ks := make([]*datastore.Key, 0, len(es))
	for i := range es {
		ks = append(ks, datastore.NewKey(c, s.kind, es[i].GetID(), 0, nil))
	}
	err := datastore.GetMulti(c, ks, es)
	if err != nil {
		if me, ok := err.(appengine.MultiError); ok {
			for i := range me {
				if s.ignoreFieldMismatch && IsErrFieldMismatch(me[i]) {
					me[i] = nil
				}
			}
		}
		return errors.Wrapf(err, "dragoon: failed to get entities - keys = %#v", ks)
	}
	return nil
}

// Put saves the entity src into the datastore based on e's ID.
func (s *Spear) Put(c context.Context, e Identifier) error {
	if err := s.CheckID(c, e); err != nil {
		return errors.Wrap(err, "dragoon: failed to generate ID")
	}
	if ts, ok := e.(TimeStamper); ok {
		SetTimeStamps(ts, Now())
	}
	if err := s.validator.Struct(e); err != nil {
		return errors.Wrap(err, "dragoon: invalid validation")
	}
	k := datastore.NewKey(c, s.kind, e.GetID(), 0, nil)
	if _, err := datastore.Put(c, k, e); err != nil {
		return errors.Wrapf(err, "dragoon: failed to put an entity - key = %#v, entity = %#v", k, e)
	}
	return nil
}

// PutMulti is a batch version of Put.
func (s *Spear) PutMulti(c context.Context, es []Identifier) error {
	now := Now()
	ks := make([]*datastore.Key, 0, len(es))
	for i := range es {
		if err := s.CheckID(c, es[i]); err != nil {
			return errors.Wrap(err, "dragoon: failed to generate new ID")
		}
		if ts, ok := es[i].(TimeStamper); ok {
			SetTimeStamps(ts, now)
		}
		if err := s.validator.Struct(es[i]); err != nil {
			return errors.Wrap(err, "dragoon: invalid validation")
		}
		ks = append(ks, datastore.NewKey(c, s.kind, es[i].GetID(), 0, nil))
	}
	if _, err := datastore.PutMulti(c, ks, es); err != nil {
		return errors.Wrapf(err, "dragoon: failed to put entities - keys = %#v, entities = %#v", ks, es)
	}
	return nil
}

// Delete deletes the entity for the given Identifier.
func (s *Spear) Delete(c context.Context, e Identifier) error {
	k := datastore.NewKey(c, s.kind, e.GetID(), 0, nil)
	if err := datastore.Delete(c, k); err != nil {
		return errors.Wrapf(err, "dragoon: failed to delete an entity - key = %#v", k)
	}
	return nil
}

// DeleteMulti is a batch version of Delete.
func (s *Spear) DeleteMulti(c context.Context, es []Identifier) error {
	ks := make([]*datastore.Key, 0, len(es))
	for i := range es {
		ks = append(ks, datastore.NewKey(c, s.kind, es[i].GetID(), 0, nil))
	}
	if err := datastore.DeleteMulti(c, ks); err != nil {
		return errors.Wrapf(err, "dragoon: failed to delete entities - keys = %#v", ks)
	}
	return nil
}

// Save saves the entity src into the datastore based on e's ID after checks exist an entity based e's ID.
func (s *Spear) Save(c context.Context, e Identifier) error {
	if err := s.CheckID(c, e); err != nil {
		return errors.Wrap(err, "dragoon: failed to generate ID")
	}
	if ts, ok := e.(TimeStamper); ok {
		SetTimeStamps(ts, Now())
	}
	if err := s.validator.Struct(e); err != nil {
		return errors.Wrap(err, "dragoon: invalid validation")
	}
	return datastore.RunInTransaction(c, func(tc context.Context) error {
		k := datastore.NewKey(tc, s.kind, e.GetID(), 0, nil)
		err := datastore.Get(tc, k, e)
		switch err {
		case nil:
			return ErrConflictEntity
		case datastore.ErrNoSuchEntity:
			_, err = datastore.Put(tc, k, e)
		}
		return err
	}, nil)
}

// CheckID checks e's ID. if e's ID is empty, set generated new ID.
func (s *Spear) CheckID(c context.Context, e Identifier) error {
	id := e.GetID()
	if id != "" {
		e.SetID(id)
		return nil
	}
	newID, err := s.identifyGenerator.NextID(c, s.kind)
	if err != nil {
		return err
	}
	e.SetID(newID)
	return nil
}

// Now returns current mills time by UTC.
func Now() time.Time {
	return time.Now().UTC().Truncate(time.Millisecond)
}

// SetTimeStamps sets a time to TimeStamper.
func SetTimeStamps(ts TimeStamper, t time.Time) {
	if ts.GetCreatedAt().IsZero() {
		ts.SetCreatedAt(t)
	}
	ts.SetUpdatedAt(t)
}

// IsErrFieldMismatch checks a type of datastore.ErrFieldMismatch or not.
func IsErrFieldMismatch(err error) bool {
	_, ok := err.(*datastore.ErrFieldMismatch)
	return ok
}

// IsNotFound checks it's datastore.ErrNoSuchEntity or not.
func IsNotFound(err error) bool {
	return err == datastore.ErrNoSuchEntity
}

// FillID fills es's ID fields.
func FillID(ks []*datastore.Key, es []Identifier) {
	for i := range ks {
		if ks[i] == nil {
			continue
		}
		es[i].SetID(ks[i].StringID())
	}
}

// AsMap converts a slice of identifiers to map.
func AsMap(is []Identifier) map[string]interface{} {
	m := make(map[string]interface{}, len(is))
	for i := range is {
		m[is[i].GetID()] = is[i]
	}
	return m
}
<commit_msg>Add Kind type<commit_after>package dragoon

import (
	"time"

	"github.com/pkg/errors"
	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
)

type (
	// Identifier gives getter and setter for ID.
	Identifier interface {
		GetID() string
		SetID(string)
	}
	// TimeStamper gives getter for CreatedAt, setter for CreatedAt and UpdatedAt.
	TimeStamper interface {
		GetCreatedAt() time.Time
		SetCreatedAt(time.Time)
		SetUpdatedAt(time.Time)
	}
	// IdentifyGenerator gives generate of ID method.
	IdentifyGenerator interface {
		NextID(c context.Context, kind string) (string, error)
	}
	// Validator gives validate of fields method.
	Validator interface {
		Struct(target interface{}) error
	}
	// Spear has convenience methods.
	Spear struct {
		kind                Kind
		ignoreFieldMismatch bool
		identifyGenerator   IdentifyGenerator
		validator           Validator
	}
	// Kind is a kind of Cloud Datastore.
	Kind string
)

// ErrConflictEntity is returned when an entity was conflict for a given key.
var ErrConflictEntity = errors.New("dragoon: conflict entity")

// NewSpear returns new Spear.
func NewSpear(k Kind, ignoreFieldMismatch bool, i IdentifyGenerator, v Validator) (*Spear, error) {
	if k == "" || i == nil || v == nil {
		return nil, errors.New("dragoon: invalid arguments")
	}
	return &Spear{
		kind:                k,
		ignoreFieldMismatch: ignoreFieldMismatch,
		identifyGenerator:   i,
		validator:           v,
	}, nil
}

// Get loads the entity based on e's key into e.
func (s *Spear) Get(c context.Context, e Identifier) error {
	k := datastore.NewKey(c, string(s.kind), e.GetID(), 0, nil)
	err := datastore.Get(c, k, e)
	if err != nil {
		if s.ignoreFieldMismatch && IsErrFieldMismatch(err) {
			return nil
		}
		return errors.Wrapf(err, "dragoon: failed to get an entity - key = %#v", k)
	}
	return nil
}

// GetMulti is a batch version of Get.
func (s *Spear) GetMulti(c context.Context, es []Identifier) error {
	ks := make([]*datastore.Key, 0, len(es))
	for i := range es {
		ks = append(ks, datastore.NewKey(c, string(s.kind), es[i].GetID(), 0, nil))
	}
	err := datastore.GetMulti(c, ks, es)
	if err != nil {
		if me, ok := err.(appengine.MultiError); ok {
			for i := range me {
				if s.ignoreFieldMismatch && IsErrFieldMismatch(me[i]) {
					me[i] = nil
				}
			}
		}
		return errors.Wrapf(err, "dragoon: failed to get entities - keys = %#v", ks)
	}
	return nil
}

// Put saves the entity src into the datastore based on e's ID.
func (s *Spear) Put(c context.Context, e Identifier) error {
	if err := s.CheckID(c, e); err != nil {
		return errors.Wrap(err, "dragoon: failed to generate ID")
	}
	if ts, ok := e.(TimeStamper); ok {
		SetTimeStamps(ts, Now())
	}
	if err := s.validator.Struct(e); err != nil {
		return errors.Wrap(err, "dragoon: invalid validation")
	}
	k := datastore.NewKey(c, string(s.kind), e.GetID(), 0, nil)
	if _, err := datastore.Put(c, k, e); err != nil {
		return errors.Wrapf(err, "dragoon: failed to put an entity - key = %#v, entity = %#v", k, e)
	}
	return nil
}

// PutMulti is a batch version of Put.
func (s *Spear) PutMulti(c context.Context, es []Identifier) error {
	now := Now()
	ks := make([]*datastore.Key, 0, len(es))
	for i := range es {
		if err := s.CheckID(c, es[i]); err != nil {
			return errors.Wrap(err, "dragoon: failed to generate new ID")
		}
		if ts, ok := es[i].(TimeStamper); ok {
			SetTimeStamps(ts, now)
		}
		if err := s.validator.Struct(es[i]); err != nil {
			return errors.Wrap(err, "dragoon: invalid validation")
		}
		ks = append(ks, datastore.NewKey(c, string(s.kind), es[i].GetID(), 0, nil))
	}
	if _, err := datastore.PutMulti(c, ks, es); err != nil {
		return errors.Wrapf(err, "dragoon: failed to put entities - keys = %#v, entities = %#v", ks, es)
	}
	return nil
}

// Delete deletes the entity for the given Identifier.
func (s *Spear) Delete(c context.Context, e Identifier) error {
	k := datastore.NewKey(c, string(s.kind), e.GetID(), 0, nil)
	if err := datastore.Delete(c, k); err != nil {
		return errors.Wrapf(err, "dragoon: failed to delete an entity - key = %#v", k)
	}
	return nil
}

// DeleteMulti is a batch version of Delete.
func (s *Spear) DeleteMulti(c context.Context, es []Identifier) error {
	ks := make([]*datastore.Key, 0, len(es))
	for i := range es {
		ks = append(ks, datastore.NewKey(c, string(s.kind), es[i].GetID(), 0, nil))
	}
	if err := datastore.DeleteMulti(c, ks); err != nil {
		return errors.Wrapf(err, "dragoon: failed to delete entities - keys = %#v", ks)
	}
	return nil
}

// Save saves the entity src into the datastore based on e's ID after checks exist an entity based e's ID.
func (s *Spear) Save(c context.Context, e Identifier) error {
	if err := s.CheckID(c, e); err != nil {
		return errors.Wrap(err, "dragoon: failed to generate ID")
	}
	if ts, ok := e.(TimeStamper); ok {
		SetTimeStamps(ts, Now())
	}
	if err := s.validator.Struct(e); err != nil {
		return errors.Wrap(err, "dragoon: invalid validation")
	}
	return datastore.RunInTransaction(c, func(tc context.Context) error {
		k := datastore.NewKey(tc, string(s.kind), e.GetID(), 0, nil)
		err := datastore.Get(tc, k, e)
		switch err {
		case nil:
			return ErrConflictEntity
		case datastore.ErrNoSuchEntity:
			_, err = datastore.Put(tc, k, e)
		}
		return err
	}, nil)
}

// CheckID checks e's ID. if e's ID is empty, set generated new ID.
func (s *Spear) CheckID(c context.Context, e Identifier) error {
	id := e.GetID()
	if id != "" {
		e.SetID(id)
		return nil
	}
	newID, err := s.identifyGenerator.NextID(c, string(s.kind))
	if err != nil {
		return err
	}
	e.SetID(newID)
	return nil
}

// Now returns current mills time by UTC.
func Now() time.Time {
	return time.Now().UTC().Truncate(time.Millisecond)
}

// SetTimeStamps sets a time to TimeStamper.
func SetTimeStamps(ts TimeStamper, t time.Time) {
	if ts.GetCreatedAt().IsZero() {
		ts.SetCreatedAt(t)
	}
	ts.SetUpdatedAt(t)
}

// IsErrFieldMismatch checks a type of datastore.ErrFieldMismatch or not.
func IsErrFieldMismatch(err error) bool {
	_, ok := err.(*datastore.ErrFieldMismatch)
	return ok
}

// IsNotFound checks it's datastore.ErrNoSuchEntity or not.
func IsNotFound(err error) bool {
	return err == datastore.ErrNoSuchEntity
}

// FillID fills es's ID fields.
func FillID(ks []*datastore.Key, es []Identifier) {
	for i := range ks {
		if ks[i] == nil {
			continue
		}
		es[i].SetID(ks[i].StringID())
	}
}

// AsMap converts a slice of identifiers to map.
func AsMap(is []Identifier) map[string]interface{} {
	m := make(map[string]interface{}, len(is))
	for i := range is {
		m[is[i].GetID()] = is[i]
	}
	return m
}
<|endoftext|>


<commit_before>package xhyve

var (
	// Version should be updated by hand at each release
	Version = "0.3.2"

	// GitCommit will be overwritten automatically by the build system
	GitCommit = "HEAD"
)
<commit_msg>xhyve: bump version 0.3.3<commit_after>package xhyve

var (
	// Version should be updated by hand at each release
	Version = "0.3.3"

	// GitCommit will be overwritten automatically by the build system
	GitCommit = "HEAD"
)
<|endoftext|>


<commit_before>package golibxml

import "testing"
import "syscall"

type ElementTypeTestCase struct {
	got ElementType
	expected string
}

var element_type_tests[] ElementTypeTestCase = []ElementTypeTestCase{
	{ XML_ELEMENT_NODE, "Node" },
	{ XML_ATTRIBUTE_NODE, "Attribute" },
	{ XML_TEXT_NODE, "Text" },
}

func getRSS() uint64 {
	rusage := &syscall.Rusage{}
	ret := syscall.Getrusage(0, rusage)
	if ret == nil && rusage.Maxrss > 0 {
		return uint64(rusage.Maxrss)
	}
	return 0
}

//
// Buffer tests
func testNewBuffer(t *testing.T) (buffer *Buffer) {
	buffer = NewBuffer()
	if buffer.Ptr == nil {
		t.Fail()
	}
	return
}

func testBufferFree(t *testing.T, buffer *Buffer) {
	buffer.Free()
	if buffer.Ptr != nil {
		t.Fail()
	}
}

func TestNewBuffer(t *testing.T) {
	testNewBuffer(t)
}

func TestNewBufferLeak(t *testing.T) {
	var buffer *Buffer
	for i := 0; i < 1000000; i++ {
		buffer = testNewBuffer(t)
		buffer.Free()
	}
	if getRSS() > 4000 {
		t.Fatal("Memory leak")
	}
}

func TestNewBufferSize(t *testing.T) {
	buffer := NewBufferSize(10)
	if buffer.Ptr == nil {
		t.Fail()
	}
	return
}

func TestNewBufferSizeLeak(t *testing.T) {
	var buffer *Buffer
	for i := 0; i < 1000000; i++ {
		buffer = NewBufferSize(1024)
		buffer.Free()
	}
	if getRSS() > 4000 {
		t.Fatal("Memory leak")
	}
}

func TestBufferFree(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
}

func TestBufferWriteChar(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.WriteChar("test")
}

func TestBufferEmpty(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.WriteChar("test")
	buffer.Empty()
	if buffer.Content() != "" {
		t.Fail()
	}
}

func TestBufferResize(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	if buffer.Resize(10) != 1 {
		t.Fail()
	}
}

func TestBufferLength(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	if buffer.Length() != 0 {
		t.Fail()
	}
	buffer.WriteChar("test")
	if buffer.Length() != 4 {
		t.Fail()
	}
}

/*
func TestBufferGrow(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.Grow(128)
	if buffer.Length() != 128 {
		t.Fail()
	}
}

func TestBufferShrink(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.Shrink(128)
	if buffer.Length() != 128 {
		t.Fail()
	}
}
*/
func TestBufferCat(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.Cat("test")
}

func TestBufferContent(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.WriteChar("test")
	if buffer.Content() != "test" {
		t.Fail()
	}
}

func TestNewDoc(t *testing.T) {
	doc := NewDoc("1.0")
	defer doc.Free()
	result := doc.String()
	if result != "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n" {
		t.Fail()
	}
}

func TestNewNode(t *testing.T) {
	doc := NewDoc("1.0")
	defer doc.Free()
	node := NewNode(nil, "div")
	doc.AddChild(node)
	result := node.String()
	if result != "<div/>" {
		t.Fail()
	}
}

func TestNewNodePath(t *testing.T) {
	doc := NewDoc("1.0")
	defer doc.Free()
	node := NewNode(nil, "div")
	doc.AddChild(node)
	result := node.Path()
	if result != "/div" {
		t.Fatal("Expected: /div Got:", node.Path())
	}
}

func TestNewComment(t *testing.T) {
	doc := NewDoc("1.0")
	defer doc.Free()
	comment := doc.NewComment("this is a comment")
	doc.AddChild(comment)
	result := comment.String()
	if result != "<!--this is a comment-->" {
		t.Fail()
	}
}

func TestElementTypeString(t *testing.T) {
	for _, test := range element_type_tests {
		if test.got.String() != test.expected {
			t.Fatal("Testing node type:", test.got, "got:", test.got.String(), "expected:", test.expected)
		}
	}
}
<commit_msg>Memory footprint increased<commit_after>package golibxml

import "testing"
import "syscall"

type ElementTypeTestCase struct {
	got ElementType
	expected string
}

var element_type_tests[] ElementTypeTestCase = []ElementTypeTestCase{
	{ XML_ELEMENT_NODE, "Node" },
	{ XML_ATTRIBUTE_NODE, "Attribute" },
	{ XML_TEXT_NODE, "Text" },
}

func getRSS() uint64 {
	rusage := &syscall.Rusage{}
	ret := syscall.Getrusage(0, rusage)
	if ret == nil && rusage.Maxrss > 0 {
		return uint64(rusage.Maxrss)
	}
	return 0
}

//
// Buffer tests
func testNewBuffer(t *testing.T) (buffer *Buffer) {
	buffer = NewBuffer()
	if buffer.Ptr == nil {
		t.Fail()
	}
	return
}

func testBufferFree(t *testing.T, buffer *Buffer) {
	buffer.Free()
	if buffer.Ptr != nil {
		t.Fail()
	}
}

func TestNewBuffer(t *testing.T) {
	testNewBuffer(t)
}

func TestNewBufferLeak(t *testing.T) {
	var buffer *Buffer
	for i := 0; i < 1000000; i++ {
		buffer = testNewBuffer(t)
		buffer.Free()
	}
	if getRSS() > 5000 {
		t.Fatal("Memory leak")
	}
}

func TestNewBufferSize(t *testing.T) {
	buffer := NewBufferSize(10)
	if buffer.Ptr == nil {
		t.Fail()
	}
	return
}

func TestNewBufferSizeLeak(t *testing.T) {
	var buffer *Buffer
	for i := 0; i < 1000000; i++ {
		buffer = NewBufferSize(1024)
		buffer.Free()
	}
	if getRSS() > 5000 {
		t.Fatal("Memory leak")
	}
}

func TestBufferFree(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
}

func TestBufferWriteChar(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.WriteChar("test")
}

func TestBufferEmpty(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.WriteChar("test")
	buffer.Empty()
	if buffer.Content() != "" {
		t.Fail()
	}
}

func TestBufferResize(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	if buffer.Resize(10) != 1 {
		t.Fail()
	}
}

func TestBufferLength(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	if buffer.Length() != 0 {
		t.Fail()
	}
	buffer.WriteChar("test")
	if buffer.Length() != 4 {
		t.Fail()
	}
}

/*
func TestBufferGrow(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.Grow(128)
	if buffer.Length() != 128 {
		t.Fail()
	}
}

func TestBufferShrink(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.Shrink(128)
	if buffer.Length() != 128 {
		t.Fail()
	}
}
*/
func TestBufferCat(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.Cat("test")
}

func TestBufferContent(t *testing.T) {
	buffer := testNewBuffer(t)
	defer testBufferFree(t, buffer)
	buffer.WriteChar("test")
	if buffer.Content() != "test" {
		t.Fail()
	}
}

func TestNewDoc(t *testing.T) {
	doc := NewDoc("1.0")
	defer doc.Free()
	result := doc.String()
	if result != "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n" {
		t.Fail()
	}
}

func TestNewNode(t *testing.T) {
	doc := NewDoc("1.0")
	defer doc.Free()
	node := NewNode(nil, "div")
	doc.AddChild(node)
	result := node.String()
	if result != "<div/>" {
		t.Fail()
	}
}

func TestNewNodePath(t *testing.T) {
	doc := NewDoc("1.0")
	defer doc.Free()
	node := NewNode(nil, "div")
	doc.AddChild(node)
	result := node.Path()
	if result != "/div" {
		t.Fatal("Expected: /div Got:", node.Path())
	}
}

func TestNewComment(t *testing.T) {
	doc := NewDoc("1.0")
	defer doc.Free()
	comment := doc.NewComment("this is a comment")
	doc.AddChild(comment)
	result := comment.String()
	if result != "<!--this is a comment-->" {
		t.Fail()
	}
}

func TestElementTypeString(t *testing.T) {
	for _, test := range element_type_tests {
		if test.got.String() != test.expected {
			t.Fatal("Testing node type:", test.got, "got:", test.got.String(), "expected:", test.expected)
		}
	}
}
<|endoftext|>


<commit_before>//Copyright 2017 Huawei Technologies Co., Ltd
//
//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 govern

import (
	"github.com/ServiceComb/service-center/pkg/util"
	apt "github.com/ServiceComb/service-center/server/core"
	"github.com/ServiceComb/service-center/server/core/backend/store"
	pb "github.com/ServiceComb/service-center/server/core/proto"
	scerr "github.com/ServiceComb/service-center/server/error"
	"github.com/ServiceComb/service-center/server/infra/registry"
	serviceUtil "github.com/ServiceComb/service-center/server/service/util"
	"golang.org/x/net/context"
)

var GovernServiceAPI pb.GovernServiceCtrlServerEx = &GovernService{}

type GovernService struct {
}

func (governService *GovernService) GetServicesInfo(ctx context.Context, in *pb.GetServicesInfoRequest) (*pb.GetServicesInfoResponse, error) {
	optionMap := make(map[string]struct{}, len(in.Options))
	for _, opt := range in.Options {
		optionMap[opt] = struct{}{}
	}

	options := make([]string, 0, len(optionMap))
	if _, ok := optionMap["all"]; ok {
		optionMap["statistics"] = struct{}{}
		options = []string{"tags", "rules", "instances", "schemas", "dependencies"}
	} else {
		for opt := range optionMap {
			options = append(options, opt)
		}
	}

	var st *pb.Statistics
	if _, ok := optionMap["statistics"]; ok {
		var err error
		st, err = statistics(ctx)
		if err != nil {
			return &pb.GetServicesInfoResponse{
				Response: pb.CreateResponse(scerr.ErrInternal, "Statistics failed."),
			}, err
		}
		if len(optionMap) == 1 {
			return &pb.GetServicesInfoResponse{
				Response:   pb.CreateResponse(pb.Response_SUCCESS, "Statistics successfully."),
				Statistics: st,
			}, nil
		}
	}

	//获取所有服务
	services, err := serviceUtil.GetAllServiceUtil(ctx)
	if err != nil {
		util.Logger().Errorf(err, "Get all services for govern service faild.")
		return &pb.GetServicesInfoResponse{
			Response: pb.CreateResponse(scerr.ErrInternal, "Get all service failed."),
		}, err
	}

	allServiceDetails := []*pb.ServiceDetail{}
	domainProject := util.ParseDomainProject(ctx)
	for _, service := range services {
		if apt.Service.ServiceId == service.ServiceId {
			continue
		}
		if len(in.AppId) > 0 {
			if in.AppId != service.AppId {
				continue
			}
			if len(in.ServiceName) > 0 && in.ServiceName != service.ServiceName {
				continue
			}
		}

		serviceDetail, err := getServiceDetailUtil(ctx, options, domainProject, service.ServiceId, service)
		if err != nil {
			return &pb.GetServicesInfoResponse{
				Response: pb.CreateResponse(scerr.ErrInternal, "Get one service detail failed."),
			}, err
		}
		serviceDetail.MicroService = service
		allServiceDetails = append(allServiceDetails, serviceDetail)
	}

	return &pb.GetServicesInfoResponse{
		Response:          pb.CreateResponse(pb.Response_SUCCESS, "Get services info successfully."),
		AllServicesDetail: allServiceDetails,
		Statistics:        st,
	}, nil
}

func (governService *GovernService) GetServiceDetail(ctx context.Context, in *pb.GetServiceRequest) (*pb.GetServiceDetailResponse, error) {
	domainProject := util.ParseDomainProject(ctx)
	options := []string{"tags", "rules", "instances", "schemas", "dependencies"}

	if len(in.ServiceId) == 0 {
		return &pb.GetServiceDetailResponse{
			Response: pb.CreateResponse(scerr.ErrInvalidParams, "Invalid requtest for getting service detail."),
		}, nil
	}

	service, err := serviceUtil.GetService(ctx, domainProject, in.ServiceId)
	if service == nil {
		return &pb.GetServiceDetailResponse{
			Response: pb.CreateResponse(scerr.ErrServiceNotExists, "Service does not exist."),
		}, nil
	}
	if err != nil {
		return &pb.GetServiceDetailResponse{
			Response: pb.CreateResponse(scerr.ErrInternal, "Get service failed."),
		}, err
	}

	key := &pb.MicroServiceKey{
		Tenant:      domainProject,
		Environment: service.Environment,
		AppId:       service.AppId,
		ServiceName: service.ServiceName,
		Version:     "",
	}
	versions, err := getServiceAllVersions(ctx, key)
	if err != nil {
		util.Logger().Errorf(err, "Get service all version fialed.")
		return &pb.GetServiceDetailResponse{
			Response: pb.CreateResponse(scerr.ErrInternal, "Get all versions of the service failed."),
		}, err
	}

	serviceInfo, err := getServiceDetailUtil(ctx, options, domainProject, in.ServiceId, service)
	if err != nil {
		return &pb.GetServiceDetailResponse{
			Response: pb.CreateResponse(scerr.ErrInternal, "Get service detail failed."),
		}, err
	}

	serviceInfo.MicroService = service
	serviceInfo.MicroServiceVersions = versions
	return &pb.GetServiceDetailResponse{
		Response: pb.CreateResponse(pb.Response_SUCCESS, "Get service successfully."),
		Service:  serviceInfo,
	}, nil
}

func (governService *GovernService) GetApplications(ctx context.Context, in *pb.GetAppsRequest) (*pb.GetAppsResponse, error) {
	err := apt.Validate(in)
	if err != nil {
		return &pb.GetAppsResponse{
			Response: pb.CreateResponse(scerr.ErrInvalidParams, err.Error()),
		}, nil
	}

	domainProject := util.ParseDomainProject(ctx)
	key := util.StringJoin([]string{
		apt.GetServiceIndexRootKey(domainProject),
		in.Environment,
	}, "/")
	if key[len(key)-1:] != "/" {
		key += "/"
	}

	opts := append(serviceUtil.FromContext(ctx),
		registry.WithStrKey(key),
		registry.WithPrefix(),
		registry.WithKeyOnly())

	resp, err := store.Store().ServiceIndex().Search(ctx, opts...)
	if err != nil {
		return nil, err
	}
	l := len(resp.Kvs)
	if l == 0 {
		return &pb.GetAppsResponse{
			Response: pb.CreateResponse(pb.Response_SUCCESS, "Get all applications successfully."),
		}, nil
	}

	apps := make([]string, 0, l)
	appMap := make(map[string]struct{}, l)
	for _, kv := range resp.Kvs {
		key, _ := pb.GetInfoFromSvcIndexKV(kv)
		if _, ok := appMap[key.AppId]; ok {
			continue
		}
		if apt.IsSCKey(key) {
			continue
		}
		appMap[key.AppId] = struct{}{}
		apps = append(apps, key.AppId)
	}

	return &pb.GetAppsResponse{
		Response: pb.CreateResponse(pb.Response_SUCCESS, "Get all applications successfully."),
		AppIds:   apps,
	}, nil
}

func getServiceAllVersions(ctx context.Context, serviceKey *pb.MicroServiceKey) ([]string, error) {
	versions := []string{}
	key := apt.GenerateServiceIndexKey(serviceKey)

	opts := append(serviceUtil.FromContext(ctx),
		registry.WithStrKey(key),
		registry.WithPrefix())

	resp, err := store.Store().ServiceIndex().Search(ctx, opts...)
	if err != nil {
		return nil, err
	}
	if resp == nil || len(resp.Kvs) == 0 {
		return versions, nil
	}
	for _, kv := range resp.Kvs {
		key, _ := pb.GetInfoFromSvcIndexKV(kv)
		versions = append(versions, key.Version)
	}
	return versions, nil
}

func getSchemaInfoUtil(ctx context.Context, domainProject string, serviceId string) ([]*pb.Schema, error) {
	key := apt.GenerateServiceSchemaKey(domainProject, serviceId, "")
	schemas := []*pb.Schema{}
	resp, err := store.Store().Schema().Search(ctx,
		registry.WithStrKey(key),
		registry.WithPrefix())
	if err != nil {
		util.Logger().Errorf(err, "Get schema failded,%s")
		return schemas, err
	}
	for _, kv := range resp.Kvs {
		schemaInfo := &pb.Schema{}
		schemaInfo.Schema = util.BytesToStringWithNoCopy(kv.Value)
		schemaInfo.SchemaId = util.BytesToStringWithNoCopy(kv.Key[len(key):])
		schemas = append(schemas, schemaInfo)
	}
	return schemas, nil
}

func getServiceDetailUtil(ctx context.Context, options []string, domainProject string, serviceId string, service *pb.MicroService) (*pb.ServiceDetail, error) {
	serviceDetail := &pb.ServiceDetail{}
	for _, opt := range options {
		expr := opt
		switch expr {
		case "tags":
			tags, err := serviceUtil.GetTagsUtils(ctx, domainProject, serviceId)
			if err != nil {
				util.Logger().Errorf(err, "Get all tags for govern service faild.")
				return nil, err
			}
			serviceDetail.Tags = tags
		case "rules":
			rules, err := serviceUtil.GetRulesUtil(ctx, domainProject, serviceId)
			if err != nil {
				util.Logger().Errorf(err, "Get all rules for govern service faild.")
				return nil, err
			}
			for _, rule := range rules {
				rule.Timestamp = rule.ModTimestamp
			}
			serviceDetail.Rules = rules
		case "instances":
			instances, err := serviceUtil.GetAllInstancesOfOneService(ctx, domainProject, serviceId)
			if err != nil {
				util.Logger().Errorf(err, "Get service's all instances for govern service faild.")
				return nil, err
			}
			serviceDetail.Instances = instances
		case "schemas":
			schemas, err := getSchemaInfoUtil(ctx, domainProject, serviceId)
			if err != nil {
				util.Logger().Errorf(err, "Get service's all schemas for govern service faild.")
				return nil, err
			}
			serviceDetail.SchemaInfos = schemas
		case "dependencies":
			dr := serviceUtil.NewDependencyRelation(ctx, domainProject, serviceId, service, serviceId, service)
			consumers, err := dr.GetDependencyConsumers()
			if err != nil {
				util.Logger().Errorf(err, "Get service's all consumers for govern service faild.")
				return nil, err
			}
			consumers = skipSelfDependency(consumers, serviceId)

			providers, err := dr.GetDependencyProviders()
			if err != nil {
				util.Logger().Errorf(err, "Get service's all providers for govern service faild.")
				return nil, err
			}
			providers = skipSelfDependency(providers, serviceId)
			serviceDetail.Consumers = consumers
			serviceDetail.Providers = providers
		case "":
			continue
		default:
			util.Logger().Errorf(nil, "option %s from request is invalid.", opt)
		}
	}
	return serviceDetail, nil
}

func skipSelfDependency(services []*pb.MicroService, serviceId string) []*pb.MicroService {
	for key, service := range services {
		if service.ServiceId == serviceId {
			services = append(services[:key], services[key+1:]...)
		}
	}
	return services
}

func statistics(ctx context.Context) (*pb.Statistics, error) {
	result := &pb.Statistics{
		Services:  &pb.StService{},
		Instances: &pb.StInstance{},
		Apps:      &pb.StApp{},
	}
	domainProject := util.ParseDomainProject(ctx)
	opts := serviceUtil.FromContext(ctx)

	// services
	key := apt.GetServiceIndexRootKey(domainProject)
	svcOpts := append(opts,
		registry.WithStrKey(key),
		registry.WithPrefix(),
		registry.WithKeyOnly())
	respSvc, err := store.Store().ServiceIndex().Search(ctx, svcOpts...)
	if err != nil {
		return nil, err
	}

	app := make(map[string]interface{}, respSvc.Count)
	scSvc := make(map[string]interface{}, respSvc.Count)
	for _, kv := range respSvc.Kvs {
		key, _ := pb.GetInfoFromSvcIndexKV(kv)
		if _, ok := app[key.AppId]; !ok {
			if !apt.IsSCKey(key) {
				app[key.AppId] = nil
			}
		}

		if apt.IsSCKey(key) {
			k := util.BytesToStringWithNoCopy(kv.Key)
			if _, ok := scSvc[k]; !ok {
				scSvc[k] = nil
			}
		}
	}
	result.Services.Count = respSvc.Count - int64(len(scSvc))
	result.Apps.Count = int64(len(app))

	// instance
	key = apt.GetInstanceRootKey(domainProject)
	instOpts := append(opts,
		registry.WithStrKey(key),
		registry.WithPrefix(),
		registry.WithKeyOnly())
	respIns, err := store.Store().Instance().Search(ctx, instOpts...)
	if err != nil {
		return nil, err
	}

	onlineServices := make(map[string]interface{}, respSvc.Count)
	for _, kv := range respIns.Kvs {
		serviceId, _, _, _ := pb.GetInfoFromInstKV(kv)
		if _, ok := onlineServices[serviceId]; !ok {
			onlineServices[serviceId] = nil
		}
	}

	key = apt.GenerateInstanceKey(domainProject, apt.Service.ServiceId, "")
	scOpts := append(opts,
		registry.WithStrKey(key),
		registry.WithPrefix(),
		registry.WithCountOnly())
	respScIns, err := store.Store().Instance().Search(ctx, scOpts...)
	if err != nil {
		return nil, err
	}

	result.Instances.Count = respIns.Count - respScIns.Count
	result.Services.OnlineCount = removeSCSelf(domainProject, int64(len(onlineServices)), 1)
	return result, err
}

func removeSCSelf(domainProject string, count int64, removeNum int64) int64 {
	if count > 0 {
		if apt.IsDefaultDomainProject(domainProject) {
			count = count - removeNum
		}
	}
	return count
}
<commit_msg>service count calculation rule: no version for one service (#198)<commit_after>//Copyright 2017 Huawei Technologies Co., Ltd
//
//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 govern

import (
	"github.com/ServiceComb/service-center/pkg/util"
	apt "github.com/ServiceComb/service-center/server/core"
	"github.com/ServiceComb/service-center/server/core/backend/store"
	pb "github.com/ServiceComb/service-center/server/core/proto"
	scerr "github.com/ServiceComb/service-center/server/error"
	"github.com/ServiceComb/service-center/server/infra/registry"
	serviceUtil "github.com/ServiceComb/service-center/server/service/util"
	"golang.org/x/net/context"
)

var GovernServiceAPI pb.GovernServiceCtrlServerEx = &GovernService{}

type GovernService struct {
}

func (governService *GovernService) GetServicesInfo(ctx context.Context, in *pb.GetServicesInfoRequest) (*pb.GetServicesInfoResponse, error) {
	optionMap := make(map[string]struct{}, len(in.Options))
	for _, opt := range in.Options {
		optionMap[opt] = struct{}{}
	}

	options := make([]string, 0, len(optionMap))
	if _, ok := optionMap["all"]; ok {
		optionMap["statistics"] = struct{}{}
		options = []string{"tags", "rules", "instances", "schemas", "dependencies"}
	} else {
		for opt := range optionMap {
			options = append(options, opt)
		}
	}

	var st *pb.Statistics
	if _, ok := optionMap["statistics"]; ok {
		var err error
		st, err = statistics(ctx)
		if err != nil {
			return &pb.GetServicesInfoResponse{
				Response: pb.CreateResponse(scerr.ErrInternal, "Statistics failed."),
			}, err
		}
		if len(optionMap) == 1 {
			return &pb.GetServicesInfoResponse{
				Response:   pb.CreateResponse(pb.Response_SUCCESS, "Statistics successfully."),
				Statistics: st,
			}, nil
		}
	}

	//获取所有服务
	services, err := serviceUtil.GetAllServiceUtil(ctx)
	if err != nil {
		util.Logger().Errorf(err, "Get all services for govern service faild.")
		return &pb.GetServicesInfoResponse{
			Response: pb.CreateResponse(scerr.ErrInternal, "Get all service failed."),
		}, err
	}

	allServiceDetails := make([]*pb.ServiceDetail, 0, len(services))
	domainProject := util.ParseDomainProject(ctx)
	for _, service := range services {
		if apt.Service.ServiceId == service.ServiceId {
			continue
		}
		if len(in.AppId) > 0 {
			if in.AppId != service.AppId {
				continue
			}
			if len(in.ServiceName) > 0 && in.ServiceName != service.ServiceName {
				continue
			}
		}

		serviceDetail, err := getServiceDetailUtil(ctx, options, domainProject, service.ServiceId, service)
		if err != nil {
			return &pb.GetServicesInfoResponse{
				Response: pb.CreateResponse(scerr.ErrInternal, "Get one service detail failed."),
			}, err
		}
		serviceDetail.MicroService = service
		allServiceDetails = append(allServiceDetails, serviceDetail)
	}

	return &pb.GetServicesInfoResponse{
		Response:          pb.CreateResponse(pb.Response_SUCCESS, "Get services info successfully."),
		AllServicesDetail: allServiceDetails,
		Statistics:        st,
	}, nil
}

func (governService *GovernService) GetServiceDetail(ctx context.Context, in *pb.GetServiceRequest) (*pb.GetServiceDetailResponse, error) {
	domainProject := util.ParseDomainProject(ctx)
	options := []string{"tags", "rules", "instances", "schemas", "dependencies"}

	if len(in.ServiceId) == 0 {
		return &pb.GetServiceDetailResponse{
			Response: pb.CreateResponse(scerr.ErrInvalidParams, "Invalid requtest for getting service detail."),
		}, nil
	}

	service, err := serviceUtil.GetService(ctx, domainProject, in.ServiceId)
	if service == nil {
		return &pb.GetServiceDetailResponse{
			Response: pb.CreateResponse(scerr.ErrServiceNotExists, "Service does not exist."),
		}, nil
	}
	if err != nil {
		return &pb.GetServiceDetailResponse{
			Response: pb.CreateResponse(scerr.ErrInternal, "Get service failed."),
		}, err
	}

	key := &pb.MicroServiceKey{
		Tenant:      domainProject,
		Environment: service.Environment,
		AppId:       service.AppId,
		ServiceName: service.ServiceName,
		Version:     "",
	}
	versions, err := getServiceAllVersions(ctx, key)
	if err != nil {
		util.Logger().Errorf(err, "Get service all version fialed.")
		return &pb.GetServiceDetailResponse{
			Response: pb.CreateResponse(scerr.ErrInternal, "Get all versions of the service failed."),
		}, err
	}

	serviceInfo, err := getServiceDetailUtil(ctx, options, domainProject, in.ServiceId, service)
	if err != nil {
		return &pb.GetServiceDetailResponse{
			Response: pb.CreateResponse(scerr.ErrInternal, "Get service detail failed."),
		}, err
	}

	serviceInfo.MicroService = service
	serviceInfo.MicroServiceVersions = versions
	return &pb.GetServiceDetailResponse{
		Response: pb.CreateResponse(pb.Response_SUCCESS, "Get service successfully."),
		Service:  serviceInfo,
	}, nil
}

func (governService *GovernService) GetApplications(ctx context.Context, in *pb.GetAppsRequest) (*pb.GetAppsResponse, error) {
	err := apt.Validate(in)
	if err != nil {
		return &pb.GetAppsResponse{
			Response: pb.CreateResponse(scerr.ErrInvalidParams, err.Error()),
		}, nil
	}

	domainProject := util.ParseDomainProject(ctx)
	key := util.StringJoin([]string{
		apt.GetServiceIndexRootKey(domainProject),
		in.Environment,
	}, "/")
	if key[len(key)-1:] != "/" {
		key += "/"
	}

	opts := append(serviceUtil.FromContext(ctx),
		registry.WithStrKey(key),
		registry.WithPrefix(),
		registry.WithKeyOnly())

	resp, err := store.Store().ServiceIndex().Search(ctx, opts...)
	if err != nil {
		return nil, err
	}
	l := len(resp.Kvs)
	if l == 0 {
		return &pb.GetAppsResponse{
			Response: pb.CreateResponse(pb.Response_SUCCESS, "Get all applications successfully."),
		}, nil
	}

	apps := make([]string, 0, l)
	appMap := make(map[string]struct{}, l)
	for _, kv := range resp.Kvs {
		key, _ := pb.GetInfoFromSvcIndexKV(kv)
		if _, ok := appMap[key.AppId]; ok {
			continue
		}
		if apt.IsSCKey(key) {
			continue
		}
		appMap[key.AppId] = struct{}{}
		apps = append(apps, key.AppId)
	}

	return &pb.GetAppsResponse{
		Response: pb.CreateResponse(pb.Response_SUCCESS, "Get all applications successfully."),
		AppIds:   apps,
	}, nil
}

func getServiceAllVersions(ctx context.Context, serviceKey *pb.MicroServiceKey) ([]string, error) {
	versions := []string{}
	key := apt.GenerateServiceIndexKey(serviceKey)

	opts := append(serviceUtil.FromContext(ctx),
		registry.WithStrKey(key),
		registry.WithPrefix())

	resp, err := store.Store().ServiceIndex().Search(ctx, opts...)
	if err != nil {
		return nil, err
	}
	if resp == nil || len(resp.Kvs) == 0 {
		return versions, nil
	}
	for _, kv := range resp.Kvs {
		key, _ := pb.GetInfoFromSvcIndexKV(kv)
		versions = append(versions, key.Version)
	}
	return versions, nil
}

func getSchemaInfoUtil(ctx context.Context, domainProject string, serviceId string) ([]*pb.Schema, error) {
	key := apt.GenerateServiceSchemaKey(domainProject, serviceId, "")

	resp, err := store.Store().Schema().Search(ctx,
		registry.WithStrKey(key),
		registry.WithPrefix())
	if err != nil {
		util.Logger().Errorf(err, "Get schema failed,%s")
		return make([]*pb.Schema, 0), err
	}
	schemas := make([]*pb.Schema, 0, len(resp.Kvs))
	for _, kv := range resp.Kvs {
		schemaInfo := &pb.Schema{}
		schemaInfo.Schema = util.BytesToStringWithNoCopy(kv.Value)
		schemaInfo.SchemaId = util.BytesToStringWithNoCopy(kv.Key[len(key):])
		schemas = append(schemas, schemaInfo)
	}
	return schemas, nil
}

func getServiceDetailUtil(ctx context.Context, options []string, domainProject string, serviceId string, service *pb.MicroService) (*pb.ServiceDetail, error) {
	serviceDetail := &pb.ServiceDetail{}
	for _, opt := range options {
		expr := opt
		switch expr {
		case "tags":
			tags, err := serviceUtil.GetTagsUtils(ctx, domainProject, serviceId)
			if err != nil {
				util.Logger().Errorf(err, "Get all tags for govern service faild.")
				return nil, err
			}
			serviceDetail.Tags = tags
		case "rules":
			rules, err := serviceUtil.GetRulesUtil(ctx, domainProject, serviceId)
			if err != nil {
				util.Logger().Errorf(err, "Get all rules for govern service faild.")
				return nil, err
			}
			for _, rule := range rules {
				rule.Timestamp = rule.ModTimestamp
			}
			serviceDetail.Rules = rules
		case "instances":
			instances, err := serviceUtil.GetAllInstancesOfOneService(ctx, domainProject, serviceId)
			if err != nil {
				util.Logger().Errorf(err, "Get service's all instances for govern service faild.")
				return nil, err
			}
			serviceDetail.Instances = instances
		case "schemas":
			schemas, err := getSchemaInfoUtil(ctx, domainProject, serviceId)
			if err != nil {
				util.Logger().Errorf(err, "Get service's all schemas for govern service faild.")
				return nil, err
			}
			serviceDetail.SchemaInfos = schemas
		case "dependencies":
			dr := serviceUtil.NewDependencyRelation(ctx, domainProject, serviceId, service, serviceId, service)
			consumers, err := dr.GetDependencyConsumers()
			if err != nil {
				util.Logger().Errorf(err, "Get service's all consumers for govern service faild.")
				return nil, err
			}
			consumers = skipSelfDependency(consumers, serviceId)

			providers, err := dr.GetDependencyProviders()
			if err != nil {
				util.Logger().Errorf(err, "Get service's all providers for govern service faild.")
				return nil, err
			}
			providers = skipSelfDependency(providers, serviceId)
			serviceDetail.Consumers = consumers
			serviceDetail.Providers = providers
		case "":
			continue
		default:
			util.Logger().Errorf(nil, "option %s from request is invalid.", opt)
		}
	}
	return serviceDetail, nil
}

func skipSelfDependency(services []*pb.MicroService, serviceId string) []*pb.MicroService {
	for key, service := range services {
		if service.ServiceId == serviceId {
			services = append(services[:key], services[key+1:]...)
		}
	}
	return services
}

func statistics(ctx context.Context) (*pb.Statistics, error) {
	result := &pb.Statistics{
		Services:  &pb.StService{},
		Instances: &pb.StInstance{},
		Apps:      &pb.StApp{},
	}
	domainProject := util.ParseDomainProject(ctx)
	opts := serviceUtil.FromContext(ctx)

	// services
	key := apt.GetServiceIndexRootKey(domainProject)
	svcOpts := append(opts,
		registry.WithStrKey(key),
		registry.WithPrefix(),
		registry.WithKeyOnly())
	respSvc, err := store.Store().ServiceIndex().Search(ctx, svcOpts...)
	if err != nil {
		return nil, err
	}

	app := make(map[string]interface{}, respSvc.Count)
	scSvc := make(map[string]interface{}, respSvc.Count)
	svcWithNonVersion := make(map[string]interface{}, respSvc.Count)
	for _, kv := range respSvc.Kvs {
		key, _ := pb.GetInfoFromSvcIndexKV(kv)
		if _, ok := app[key.AppId]; !ok {
			if !apt.IsSCKey(key) {
				app[key.AppId] = nil
			}
		}

		if apt.IsSCKey(key) {
			k := util.BytesToStringWithNoCopy(kv.Key)
			if _, ok := scSvc[k]; !ok {
				scSvc[k] = nil
			}
		}

		key.Version = ""
		svcWithNonVersionKey := apt.GenerateServiceIndexKey(key)
		svcWithNonVersion[svcWithNonVersionKey] = nil
	}
	scSvcCount := 0
	if int64(len(scSvc)) > 0 {
		scSvcCount = 1
	}
	result.Services.Count = int64(len(svcWithNonVersion) - scSvcCount)
	result.Apps.Count = int64(len(app))

	// instance
	key = apt.GetInstanceRootKey(domainProject)
	instOpts := append(opts,
		registry.WithStrKey(key),
		registry.WithPrefix(),
		registry.WithKeyOnly())
	respIns, err := store.Store().Instance().Search(ctx, instOpts...)
	if err != nil {
		return nil, err
	}

	onlineServices := make(map[string]interface{}, respSvc.Count)
	for _, kv := range respIns.Kvs {
		serviceId, _, _, _ := pb.GetInfoFromInstKV(kv)
		if _, ok := onlineServices[serviceId]; !ok {
			onlineServices[serviceId] = nil
		}
	}

	key = apt.GenerateInstanceKey(domainProject, apt.Service.ServiceId, "")
	scOpts := append(opts,
		registry.WithStrKey(key),
		registry.WithPrefix(),
		registry.WithCountOnly())
	respScIns, err := store.Store().Instance().Search(ctx, scOpts...)
	if err != nil {
		return nil, err
	}

	result.Instances.Count = respIns.Count - respScIns.Count
	result.Services.OnlineCount = removeSCSelf(domainProject, int64(len(onlineServices)), 1)
	return result, err
}

func removeSCSelf(domainProject string, count int64, removeNum int64) int64 {
	if count > 0 {
		if apt.IsDefaultDomainProject(domainProject) {
			count = count - removeNum
		}
	}
	return count
}
<|endoftext|>


<commit_before>package docker

import (
	"fmt"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/mount"
	"github.com/docker/go-connections/nat"
	"github.com/go-errors/errors"
	"github.com/viant/toolbox"
	"github.com/viant/toolbox/secret"
	"github.com/viant/toolbox/url"
	"strings"
)

//RunRequest represents a docker runAdapter request
type RunRequest struct {
	Credentials                 string `description:"credentials"`
	Name                        string
	Image                       string            `required:"true" description:"container image to runAdapter" example:"mysql:5.6"`
	Port                        string            `description:"publish a container’s port(s) to the host, docker -p option"`
	Env                         map[string]string `description:"set docker container an environment variable, docker -e KEY=VAL  option"`
	Mount                       map[string]string `description:"bind mount a volume, docker -v option"`
	Ports                       map[string]string `description:"publish a container’s port(s) to the host, docker -p option"`
	Workdir                     string            `description:"working directory inside the container, docker -w option"`
	Reuse                       bool              `description:"reuse existing container if exists, otherwise always removes"`
	Cmd                         []string
	Entrypoint                  []string
	types.ContainerCreateConfig `json:",inline" yaml:",inline"`
	Secrets                     map[secret.SecretKey]secret.Secret `description:"map of secrets used within env"`
}

type RunResponse struct {
	ContainerID string
	Status      string
	Stdout      string
}

//BuildRequest represents docker build request
type BuildRequest struct {
	Tag                     *Tag   `required:"true" description:"build docker tag"`
	Path                    string `description:"location of dockerfile"`
	types.ImageBuildOptions `json:",inline" yaml:",inline"`
}

func (r *BuildRequest) Init() error {
	if r.Path == "" {
		r.Path = url.NewResource(".").ParsedURL.Path
	}
	if len(r.Tags) == 0 {
		r.Tags = make([]string, 0)
		if r.Tag != nil {
			r.Tags = append(r.Tags, r.Tag.String())
		}
	}
	return nil
}

//BuildResponse represents image ID
type BuildResponse struct {
	ImageID string
	Stdout  []string
}

//LoginRequest represents a docker pull request
type LoginRequest struct {
	Credentials string `required:"true" description:"credentials path"`
	Repository  string `required:"true" description:"repository url"`
}

//LoginResponse represents login response
type LoginResponse struct {
}

//TagRequest represents docker tag request
type TagRequest struct {
	SourceTag *Tag `required:"true"`
	TargetTag *Tag `required:"true"`
}

//TagResponse represents docker tag response
type TagResponse struct {
	Stdout string
}

//PushRequest represents a docker push request
type PushRequest struct {
	Credentials string
	Tag         *Tag `required:"true"`
}

//PushResponse represents push response
type PushResponse struct {
	Stdout []string
}

//StatusRequest represents a docker check container status request
type StatusRequest struct {
	Name   string
	Names  []string
	Images []string
	IDs    []string
}

//StatusResponse represents status response
type StatusResponse struct {
	Containers []types.Container
}

//StartRequest start request
type StartRequest StatusRequest

//StartResponse represents docker start response
type StartResponse StopResponse

//StopRequest represents docker stop running images/containers request
type StopRequest StatusRequest

//StopImagesResponse represents docker stop images response
type StopResponse StatusResponse

//RemoveRequest represents docker remove request
type RemoveRequest StatusRequest

//RemoveResponse represents remove response
type RemoveResponse StatusResponse

//LogsRequest represents docker runner container logs to take stdout
type LogsRequest struct {
	StatusRequest
	*types.ContainerLogsOptions
}

//LogsResponse represents docker container logs response
type LogsResponse struct {
	Stdout string
}

//PullRequest represents pull request
type PullRequest struct {
	Credentials            string
	Image                  string
	types.ImagePullOptions `json:",inline" yaml:",inline"`
}

//PullResponse represents pull response
type PullResponse struct {
	types.ImageSummary
	Stdout []string
}

//LogoutRequest represents a docker logout request
type LogoutRequest struct {
	Repository string `required:"true" description:"repository URL"`
}

//LogoutResponse represents a docker logout response
type LogoutResponse struct{}

//PushResponse represents a docker push request
type CopyRequest struct {
	Assets map[string]string
}

//CopyResponse represents a copy response
type CopyResponse struct{}

//InspectRequest represents a docker inspect request, target name refers to container name
type InspectRequest StatusRequest

//InspectResponse represents a docker inspect request
type InspectResponse struct {
	Info []types.ContainerJSON //you can extract any instance default, for instance to get Ip you can use Info[0].NetworkSettings.IPAddress in the variable action post from key
}

func (r *CopyRequest) Validate() error {
	if len(r.Assets) == 0 {
		return fmt.Errorf("asset was empty")
	}

	return nil
}

func (r *RunRequest) Init() error {
	if r.Config == nil {
		r.Config = &container.Config{}
	}
	if r.HostConfig == nil {
		r.HostConfig = &container.HostConfig{}
	}

	if r.Image != "" {
		r.Config.Image = r.Image
	}
	if len(r.Mount) > 0 {
		r.HostConfig.Mounts = make([]mount.Mount, 0)
		r.Config.Volumes = make(map[string]struct{})
		for source, dest := range r.Mount {
			if parts := strings.SplitN(source, ":", 2); len(parts) == 2 {
				source = parts[0]
				dest = parts[1]
			}
			source = expandHomeDirectory(source)
			source = url.NewResource(source).ParsedURL.Path
			r.HostConfig.Mounts = append(r.HostConfig.Mounts, mount.Mount{
				Type:   mount.TypeBind,
				Source: source,
				Target: dest,
			})
		}
	}
	if r.Port != "" {
		portSet := nat.PortSet{nat.Port(r.Port): struct{}{}}
		if err := toolbox.DefaultConverter.AssignConverted(&r.Config.ExposedPorts, portSet);err != nil {
			return err
		}
	}
	if len(r.Ports) > 0 {
		for source, dest := range r.Ports {
			if !strings.Contains(dest, "/") {
				dest += "/tcp"
			}
			ports := []nat.PortBinding{{HostIP: "0.0.0.0", HostPort: source}}
			var portsBindings = make(map[nat.Port][]nat.PortBinding)
			portsBindings[nat.Port(dest)] = ports
			if err := toolbox.DefaultConverter.AssignConverted(&r.HostConfig.PortBindings, portsBindings);err != nil {
				return err
			}
		}
	}
	if len(r.Env) > 0 {
		r.Config.Env = make([]string, 0)
		for k, v := range r.Env {
			r.Config.Env = append(r.Config.Env, fmt.Sprintf("%v=%v", k, v))
		}
	}
	if r.Workdir != "" {
		r.Config.WorkingDir = r.Workdir
	}
	if len(r.Cmd) > 0 {
		r.Config.Cmd = r.Cmd
	}
	if len(r.Entrypoint) > 0 {
		r.Config.Entrypoint = r.Entrypoint
	}
	if r.Name != "" {
		r.ContainerCreateConfig.Name = r.Name
	}
	return nil
}

func (r *RunRequest) Validate() error {
	if r.Config.Image == "" {
		return errors.New("image was empty")
	}
	return nil
}

func (r *PullRequest) Init() error {
	return nil
}

func (r *RunRequest) CreateContainerRequest() *ContainerCreateRequest {
	createRequest := &ContainerCreateRequest{}
	createRequest.Config = r.ContainerCreateConfig.Config
	createRequest.NetworkingConfig = r.ContainerCreateConfig.NetworkingConfig
	createRequest.HostConfig = r.ContainerCreateConfig.HostConfig
	createRequest.ContainerName = r.Name
	return createRequest
}

func (r *LoginRequest) Validate() error {
	if r.Credentials == "" {
		return errors.New("credentials were empty")
	}
	if r.Repository == "" {
		return errors.New("repository was empty")
	}
	return nil
}

func (r *StatusRequest) Init() error {
	if r.Name != "" && len(r.Names) == 0 {
		r.Names = strings.Split(r.Name, ",")
	}
	return nil
}

//StatusRequest returns status request
func (r *StopRequest) AsStatusRequest() *StatusRequest {
	result := StatusRequest(*r)
	return &result
}

//StatusRequest returns status request
func (r *RemoveRequest) AsStatusRequest() *StatusRequest {
	result := StatusRequest(*r)
	return &result
}

//StatusRequest returns status request
func (r *StartRequest) AsStatusRequest() *StatusRequest {
	result := StatusRequest(*r)
	return &result
}

//StatusRequest returns status request
func (r *InspectRequest) AsStatusRequest() *StatusRequest {
	result := StatusRequest(*r)
	return &result
}

//StatusRequest returns status request
func (r *LogsRequest) AsStatusRequest() *StatusRequest {
	return &r.StatusRequest
}

//StatusRequest returns status request
func (r *LogsRequest) Init() error {
	if r.ContainerLogsOptions == nil {
		r.ContainerLogsOptions = &types.ContainerLogsOptions{
			ShowStdout: true,
			ShowStderr: true,
		}
	}
	return nil
}
<commit_msg>patched contract<commit_after>package docker

import (
	"fmt"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/mount"
	"github.com/docker/go-connections/nat"
	"github.com/go-errors/errors"
	"github.com/viant/toolbox"
	"github.com/viant/toolbox/secret"
	"github.com/viant/toolbox/url"
	"strings"
)

//RunRequest represents a docker runAdapter request
type RunRequest struct {
	Credentials                 string `description:"credentials"`
	Name                        string
	Image                       string            `required:"true" description:"container image to runAdapter" example:"mysql:5.6"`
	Port                        string            `description:"publish a container’s port(s) to the host, docker -p option"`
	Env                         map[string]string `description:"set docker container an environment variable, docker -e KEY=VAL  option"`
	Mount                       map[string]string `description:"bind mount a volume, docker -v option"`
	Ports                       map[string]string `description:"publish a container’s port(s) to the host, docker -p option"`
	Workdir                     string            `description:"working directory inside the container, docker -w option"`
	Reuse                       bool              `description:"reuse existing container if exists, otherwise always removes"`
	Cmd                         []string
	Entrypoint                  []string
	types.ContainerCreateConfig `json:",inline" yaml:",inline"`
	Secrets                     map[secret.SecretKey]secret.Secret `description:"map of secrets used within env"`
}

type RunResponse struct {
	ContainerID string
	Status      string
	Stdout      string
}

//BuildRequest represents docker build request
type BuildRequest struct {
	Tag                     *Tag   `required:"true" description:"build docker tag"`
	Path                    string `description:"location of dockerfile"`
	types.ImageBuildOptions `json:",inline" yaml:",inline"`
}

func (r *BuildRequest) Init() error {
	if r.Path == "" {
		r.Path = url.NewResource(".").ParsedURL.Path
	}
	if len(r.Tags) == 0 {
		r.Tags = make([]string, 0)
		if r.Tag != nil {
			r.Tags = append(r.Tags, r.Tag.String())
		}
	}
	return nil
}

//BuildResponse represents image ID
type BuildResponse struct {
	ImageID string
	Stdout  []string
}

//LoginRequest represents a docker pull request
type LoginRequest struct {
	Credentials string `required:"true" description:"credentials path"`
	Repository  string `required:"true" description:"repository url"`
}

//LoginResponse represents login response
type LoginResponse struct {
}

//TagRequest represents docker tag request
type TagRequest struct {
	SourceTag *Tag `required:"true"`
	TargetTag *Tag `required:"true"`
}

//TagResponse represents docker tag response
type TagResponse struct {
	Stdout string
}

//PushRequest represents a docker push request
type PushRequest struct {
	Credentials string
	Tag         *Tag `required:"true"`
}

//PushResponse represents push response
type PushResponse struct {
	Stdout []string
}

//StatusRequest represents a docker check container status request
type StatusRequest struct {
	Name   string
	Names  []string
	Images []string
	IDs    []string
}

//StatusResponse represents status response
type StatusResponse struct {
	Containers []types.Container
}

//StartRequest start request
type StartRequest StatusRequest

//StartResponse represents docker start response
type StartResponse StopResponse

//StopRequest represents docker stop running images/containers request
type StopRequest StatusRequest

//StopImagesResponse represents docker stop images response
type StopResponse StatusResponse

//RemoveRequest represents docker remove request
type RemoveRequest StatusRequest

//RemoveResponse represents remove response
type RemoveResponse StatusResponse

//LogsRequest represents docker runner container logs to take stdout
type LogsRequest struct {
	StatusRequest
	*types.ContainerLogsOptions
}

//LogsResponse represents docker container logs response
type LogsResponse struct {
	Stdout string
}

//PullRequest represents pull request
type PullRequest struct {
	Credentials            string
	Image                  string
	types.ImagePullOptions `json:",inline" yaml:",inline"`
}

//PullResponse represents pull response
type PullResponse struct {
	types.ImageSummary
	Stdout []string
}

//LogoutRequest represents a docker logout request
type LogoutRequest struct {
	Repository string `required:"true" description:"repository URL"`
}

//LogoutResponse represents a docker logout response
type LogoutResponse struct{}

//PushResponse represents a docker push request
type CopyRequest struct {
	Assets map[string]string
}

//CopyResponse represents a copy response
type CopyResponse struct{}

//InspectRequest represents a docker inspect request, target name refers to container name
type InspectRequest StatusRequest

//InspectResponse represents a docker inspect request
type InspectResponse struct {
	Info []types.ContainerJSON //you can extract any instance default, for instance to get Ip you can use Info[0].NetworkSettings.IPAddress in the variable action post from key
}

func (r *CopyRequest) Validate() error {
	if len(r.Assets) == 0 {
		return fmt.Errorf("asset was empty")
	}

	return nil
}

func (r *RunRequest) Init() error {
	if r.Config == nil {
		r.Config = &container.Config{}
	}
	if r.HostConfig == nil {
		r.HostConfig = &container.HostConfig{}
	}

	if r.Image != "" {
		r.Config.Image = r.Image
	}
	if len(r.Mount) > 0 {
		r.HostConfig.Mounts = make([]mount.Mount, 0)
		r.Config.Volumes = make(map[string]struct{})
		for source, dest := range r.Mount {
			if parts := strings.SplitN(source, ":", 2); len(parts) == 2 {
				source = parts[0]
				dest = parts[1]
			}
			source = expandHomeDirectory(source)
			source = url.NewResource(source).ParsedURL.Path
			r.HostConfig.Mounts = append(r.HostConfig.Mounts, mount.Mount{
				Type:   mount.TypeBind,
				Source: source,
				Target: dest,
			})
		}
	}
	if r.Port != "" {
		portSet := nat.PortSet{nat.Port(r.Port): struct{}{}}
		if err := toolbox.DefaultConverter.AssignConverted(&r.Config.ExposedPorts, portSet);err != nil {
			return err
		}
	}
	if len(r.Ports) > 0 {
		portsBindings := map[nat.Port][]nat.PortBinding{}
		for source, dest := range r.Ports {
			if !strings.Contains(dest, "/") {
				dest += "/tcp"
			}
			portsBindings[nat.Port(dest)] =  []nat.PortBinding{{HostIP: "0.0.0.0", HostPort: string(source)}}
		}
		if err := toolbox.DefaultConverter.AssignConverted(&r.HostConfig.PortBindings, portsBindings);err != nil {
			return err
		}

	}
	if len(r.Env) > 0 {
		r.Config.Env = make([]string, 0)
		for k, v := range r.Env {
			r.Config.Env = append(r.Config.Env, fmt.Sprintf("%v=%v", k, v))
		}
	}
	if r.Workdir != "" {
		r.Config.WorkingDir = r.Workdir
	}
	if len(r.Cmd) > 0 {
		r.Config.Cmd = r.Cmd
	}
	if len(r.Entrypoint) > 0 {
		r.Config.Entrypoint = r.Entrypoint
	}
	if r.Name != "" {
		r.ContainerCreateConfig.Name = r.Name
	}
	return nil
}

func (r *RunRequest) Validate() error {
	if r.Config.Image == "" {
		return errors.New("image was empty")
	}
	return nil
}

func (r *PullRequest) Init() error {
	return nil
}

func (r *RunRequest) CreateContainerRequest() *ContainerCreateRequest {
	createRequest := &ContainerCreateRequest{}
	createRequest.Config = r.ContainerCreateConfig.Config
	createRequest.NetworkingConfig = r.ContainerCreateConfig.NetworkingConfig
	createRequest.HostConfig = r.ContainerCreateConfig.HostConfig
	createRequest.ContainerName = r.Name
	return createRequest
}

func (r *LoginRequest) Validate() error {
	if r.Credentials == "" {
		return errors.New("credentials were empty")
	}
	if r.Repository == "" {
		return errors.New("repository was empty")
	}
	return nil
}

func (r *StatusRequest) Init() error {
	if r.Name != "" && len(r.Names) == 0 {
		r.Names = strings.Split(r.Name, ",")
	}
	return nil
}

//StatusRequest returns status request
func (r *StopRequest) AsStatusRequest() *StatusRequest {
	result := StatusRequest(*r)
	return &result
}

//StatusRequest returns status request
func (r *RemoveRequest) AsStatusRequest() *StatusRequest {
	result := StatusRequest(*r)
	return &result
}

//StatusRequest returns status request
func (r *StartRequest) AsStatusRequest() *StatusRequest {
	result := StatusRequest(*r)
	return &result
}

//StatusRequest returns status request
func (r *InspectRequest) AsStatusRequest() *StatusRequest {
	result := StatusRequest(*r)
	return &result
}

//StatusRequest returns status request
func (r *LogsRequest) AsStatusRequest() *StatusRequest {
	return &r.StatusRequest
}

//StatusRequest returns status request
func (r *LogsRequest) Init() error {
	if r.ContainerLogsOptions == nil {
		r.ContainerLogsOptions = &types.ContainerLogsOptions{
			ShowStdout: true,
			ShowStderr: true,
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright Project Harbor 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 period

import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"

	"github.com/goharbor/harbor/src/jobservice/common/rds"
	"github.com/goharbor/harbor/src/jobservice/common/utils"
	"github.com/goharbor/harbor/src/jobservice/env"
	"github.com/goharbor/harbor/src/jobservice/job"
	"github.com/goharbor/harbor/src/jobservice/lcm"
	"github.com/goharbor/harbor/src/jobservice/tests"
	"github.com/gomodule/redigo/redis"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
)

// EnqueuerTestSuite tests functions of enqueuer
type EnqueuerTestSuite struct {
	suite.Suite

	enqueuer  *enqueuer
	namespace string
	pool      *redis.Pool
	cancel    context.CancelFunc
}

// TestEnqueuerTestSuite is entry of go test
func TestEnqueuerTestSuite(t *testing.T) {
	suite.Run(t, new(EnqueuerTestSuite))
}

// SetupSuite prepares the test suite
func (suite *EnqueuerTestSuite) SetupSuite() {
	suite.namespace = tests.GiveMeTestNamespace()
	suite.pool = tests.GiveMeRedisPool()

	ctx, cancel := context.WithCancel(context.WithValue(context.Background(), utils.NodeID, "fake_node_ID"))
	suite.cancel = cancel

	envCtx := &env.Context{
		SystemContext: ctx,
		WG:            new(sync.WaitGroup),
	}

	lcmCtl := lcm.NewController(
		envCtx,
		suite.namespace,
		suite.pool,
		func(hookURL string, change *job.StatusChange) error { return nil },
	)
	suite.enqueuer = newEnqueuer(ctx, suite.namespace, suite.pool, lcmCtl)
	suite.prepare()

	suite.enqueuer.start()
}

// TearDownSuite clears the test suite
func (suite *EnqueuerTestSuite) TearDownSuite() {
	suite.cancel()

	conn := suite.pool.Get()
	defer func() {
		_ = conn.Close()
	}()

	_ = tests.ClearAll(suite.namespace, conn)
}

// TestEnqueuer tests enqueuer
func (suite *EnqueuerTestSuite) TestEnqueuer() {
	key := rds.RedisKeyScheduled(suite.namespace)
	conn := suite.pool.Get()
	defer func() {
		if err := conn.Close(); err != nil {
			suite.NoError(err, "close redis connection")
		}
	}()

	tk := time.NewTicker(500 * time.Millisecond)
	defer tk.Stop()

	for {
		select {
		case <-tk.C:
			count, err := redis.Int(conn.Do("ZCARD", key))
			require.Nil(suite.T(), err, "count scheduled: nil error expected but got %s", err)
			if assert.Condition(suite.T(), func() (success bool) {
				return count > 0
			}, "at least one job should be scheduled for the periodic job policy") {
				return
			}
		case <-time.After(15 * time.Second):
			require.NoError(suite.T(), errors.New("timeout (15s): expect at 1 scheduled job but still get nothing"))
			return
		}
	}
}

func (suite *EnqueuerTestSuite) prepare() {
	now := time.Now()
	minute := now.Minute()

	coreSpec := fmt.Sprintf("0-59 %d * * * *", minute)

	// Prepare one
	p := &Policy{
		ID:       "fake_policy",
		JobName:  job.SampleJob,
		CronSpec: coreSpec,
	}
	rawData, err := p.Serialize()
	assert.Nil(suite.T(), err, "prepare data: nil error expected but got %s", err)
	key := rds.KeyPeriodicPolicy(suite.namespace)

	conn := suite.pool.Get()
	defer func() {
		_ = conn.Close()
	}()

	_, err = conn.Do("ZADD", key, time.Now().Unix(), rawData)
	assert.Nil(suite.T(), err, "prepare policy: nil error expected but got %s", err)
}
<commit_msg>fix[UT]:improve the UT cases of enqueuer (#11358)<commit_after>// Copyright Project Harbor 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 period

import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"

	"github.com/goharbor/harbor/src/jobservice/common/rds"
	"github.com/goharbor/harbor/src/jobservice/common/utils"
	"github.com/goharbor/harbor/src/jobservice/env"
	"github.com/goharbor/harbor/src/jobservice/job"
	"github.com/goharbor/harbor/src/jobservice/lcm"
	"github.com/goharbor/harbor/src/jobservice/tests"
	"github.com/gomodule/redigo/redis"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/suite"
)

// EnqueuerTestSuite tests functions of enqueuer
type EnqueuerTestSuite struct {
	suite.Suite

	enqueuer  *enqueuer
	namespace string
	pool      *redis.Pool
	cancel    context.CancelFunc
}

// TestEnqueuerTestSuite is entry of go test
func TestEnqueuerTestSuite(t *testing.T) {
	suite.Run(t, new(EnqueuerTestSuite))
}

// SetupSuite prepares the test suite
func (suite *EnqueuerTestSuite) SetupSuite() {
	suite.namespace = tests.GiveMeTestNamespace()
	suite.pool = tests.GiveMeRedisPool()

	ctx, cancel := context.WithCancel(context.WithValue(context.Background(), utils.NodeID, "fake_node_ID"))
	suite.cancel = cancel

	envCtx := &env.Context{
		SystemContext: ctx,
		WG:            new(sync.WaitGroup),
	}

	lcmCtl := lcm.NewController(
		envCtx,
		suite.namespace,
		suite.pool,
		func(hookURL string, change *job.StatusChange) error { return nil },
	)
	suite.enqueuer = newEnqueuer(ctx, suite.namespace, suite.pool, lcmCtl)
	suite.prepare()

	suite.enqueuer.start()
}

// TearDownSuite clears the test suite
func (suite *EnqueuerTestSuite) TearDownSuite() {
	suite.cancel()

	conn := suite.pool.Get()
	defer func() {
		_ = conn.Close()
	}()

	_ = tests.ClearAll(suite.namespace, conn)
}

// TestEnqueuer tests enqueuer
func (suite *EnqueuerTestSuite) TestEnqueuer() {
	key := rds.RedisKeyScheduled(suite.namespace)
	conn := suite.pool.Get()
	defer func() {
		if err := conn.Close(); err != nil {
			suite.NoError(err, "close redis connection")
		}
	}()

	tk := time.NewTicker(497 * time.Millisecond)
	defer tk.Stop()

	tm := time.NewTimer(15 * time.Second)
	defer tm.Stop()

	for {
		select {
		case <-tk.C:
			count, err := redis.Int(conn.Do("ZCARD", key))
			require.Nil(suite.T(), err, "count scheduled: nil error expected but got %s", err)
			if assert.Condition(suite.T(), func() (success bool) {
				return count > 0
			}, "at least one job should be scheduled for the periodic job policy") {
				return
			}
		case <-tm.C:
			require.NoError(suite.T(), errors.New("timeout (15s): expect at 1 scheduled job but still get nothing"))
			return
		}
	}
}

func (suite *EnqueuerTestSuite) prepare() {
	now := time.Now()
	minute := now.Minute()

	coreSpec := fmt.Sprintf("0-59 %d-%d * * * *", minute, minute+2)

	// Prepare one
	p := &Policy{
		ID:       "fake_policy",
		JobName:  job.SampleJob,
		CronSpec: coreSpec,
	}
	rawData, err := p.Serialize()
	assert.Nil(suite.T(), err, "prepare data: nil error expected but got %s", err)
	key := rds.KeyPeriodicPolicy(suite.namespace)

	conn := suite.pool.Get()
	defer func() {
		_ = conn.Close()
	}()

	_, err = conn.Do("ZADD", key, time.Now().Unix(), rawData)
	assert.Nil(suite.T(), err, "prepare policy: nil error expected but got %s", err)
}
<|endoftext|>