handler.go

//
// Copyright (C) 2019-2020 Vdaas.org Vald team ( kpango, rinx, kmrmt )
//
// 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
//
//    https://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 grpc provides grpc server logic
package grpc

import (
	"context"
	"fmt"
	"math"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/kpango/fuid"
	"github.com/vdaas/vald/apis/grpc/agent"
	"github.com/vdaas/vald/apis/grpc/gateway/vald"
	"github.com/vdaas/vald/apis/grpc/payload"
	"github.com/vdaas/vald/internal/errgroup"
	"github.com/vdaas/vald/internal/errors"
	"github.com/vdaas/vald/internal/info"
	"github.com/vdaas/vald/internal/log"
	"github.com/vdaas/vald/internal/net/grpc"
	"github.com/vdaas/vald/internal/net/grpc/status"
	"github.com/vdaas/vald/internal/safety"
	"github.com/vdaas/vald/pkg/gateway/vald/service"
)

type server struct {
	eg                errgroup.Group
	gateway           service.Gateway
	metadata          service.Meta
	backup            service.Backup
	timeout           time.Duration
	filter            service.Filter
	replica           int
	streamConcurrency int
}

func New(opts ...Option) vald.ValdServer {
	s := new(server)

	for _, opt := range append(defaultOpts, opts...) {
		opt(s)
	}
	return s
}

func (s *server) Exists(ctx context.Context, meta *payload.Object_ID) (*payload.Object_ID, error) {
	uuid, err := s.metadata.GetUUID(ctx, meta.GetId())
	if err != nil {
		return nil, status.WrapWithNotFound(fmt.Sprintf("Exists API meta %s's uuid not found", meta.GetId()), err, meta.GetId(), info.Get())
	}
	return &payload.Object_ID{
		Id: uuid,
	}, nil
}

func (s *server) Search(ctx context.Context, req *payload.Search_Request) (res *payload.Search_Response, err error) {
	return s.search(ctx, req.GetConfig(),
		func(ctx context.Context, ac agent.AgentClient, copts ...grpc.CallOption) (*payload.Search_Response, error) {
			return ac.Search(ctx, req, copts...)
		})
}

func (s *server) SearchByID(ctx context.Context, req *payload.Search_IDRequest) (
	res *payload.Search_Response, err error) {
	metaID := req.GetId()
	req.Id, err = s.metadata.GetUUID(ctx, metaID)
	if err != nil {
		req.Id = metaID
		log.Errorf("error at SearchByID\t%v", err)
		return nil, status.WrapWithNotFound(fmt.Sprintf("SearchByID API meta %s's uuid not found", metaID), err, req, info.Get())
	}
	return s.search(ctx, req.GetConfig(),
		func(ctx context.Context, ac agent.AgentClient, copts ...grpc.CallOption) (*payload.Search_Response, error) {
			return ac.SearchByID(ctx, req, copts...)
		})
}

func (s *server) search(ctx context.Context, cfg *payload.Search_Config,
	f func(ctx context.Context, ac agent.AgentClient, copts ...grpc.CallOption) (*payload.Search_Response, error)) (
	res *payload.Search_Response, err error) {
	maxDist := uint32(math.MaxUint32)
	num := int(cfg.GetNum())
	res = new(payload.Search_Response)
	res.Results = make([]*payload.Object_Distance, 0, s.gateway.GetAgentCount()*num)
	dch := make(chan *payload.Object_Distance, cap(res.GetResults())/2)
	eg, ectx := errgroup.New(ctx)
	var cancel context.CancelFunc
	var timeout time.Duration
	if to := cfg.GetTimeout(); to != 0 {
		timeout = time.Duration(to)
	} else {
		timeout = s.timeout
	}
	ectx, cancel = context.WithTimeout(ectx, timeout)

	eg.Go(safety.RecoverFunc(func() error {
		defer cancel()
		// cl := new(checkList)
		visited := make(map[string]bool, len(res.Results))
		mu := sync.RWMutex{}
		return s.gateway.BroadCast(ectx, func(ctx context.Context, target string, ac agent.AgentClient, copts ...grpc.CallOption) error {
			r, err := f(ctx, ac, copts...)
			if err != nil {
				log.Debug(err)
				return nil
			}
			for _, dist := range r.GetResults() {
				if dist.GetDistance() > math.Float32frombits(atomic.LoadUint32(&maxDist)) {
					return nil
				}
				id := dist.GetId()
				mu.RLock()
				if !visited[id] {
					mu.RUnlock()
					mu.Lock()
					visited[id] = true
					mu.Unlock()
					dch <- dist
				} else {
					mu.RUnlock()
				}
				// if !cl.Exists(id) {
				// 	dch <- dist
				// 	cl.Check(id)
				// }
			}
			return nil
		})
	}))
	for {
		select {
		case <-ectx.Done():
			err = eg.Wait()
			if err != nil {
				log.Error(err)
			}
			close(dch)
			if len(res.GetResults()) > num && num != 0 {
				res.Results = res.Results[:num]
			}
			uuids := make([]string, 0, len(res.Results))
			for _, r := range res.Results {
				uuids = append(uuids, r.GetId())
			}
			if s.metadata != nil {
				metas, merr := s.metadata.GetMetas(ctx, uuids...)
				if merr != nil {
					log.Error(merr)
					err = errors.Wrap(err, merr.Error())
				}
				for i, k := range metas {
					if len(k) != 0 {
						res.Results[i].Id = k
					}
				}
			}
			if s.filter != nil {
				r, ferr := s.filter.FilterSearch(ctx, res)
				if ferr == nil {
					res = r
				} else {
					err = errors.Wrap(err, ferr.Error())
				}
			}
			if err != nil {
				return res, status.WrapWithInternal(fmt.Sprintf("failed to search request %#v", cfg), err, info.Get())
			}
			return res, nil
		case dist := <-dch:
			if len(res.GetResults()) >= num &&
				dist.GetDistance() < math.Float32frombits(atomic.LoadUint32(&maxDist)) {
				atomic.StoreUint32(&maxDist, math.Float32bits(dist.GetDistance()))
			}
			switch len(res.GetResults()) {
			case 0:
				res.Results = append(res.Results, dist)
				continue
			case 1:
				if res.GetResults()[0].GetDistance() <= dist.GetDistance() {
					res.Results = append(res.Results, dist)
				} else {
					res.Results = append([]*payload.Object_Distance{dist}, res.Results[0])
				}
				continue
			}

			pos := len(res.GetResults())
			for idx := pos; idx >= 1; idx-- {
				if res.GetResults()[idx-1].GetDistance() <= dist.GetDistance() {
					pos = idx - 1
					break
				}
			}
			switch {
			case pos == len(res.GetResults()):
				res.Results = append([]*payload.Object_Distance{dist}, res.Results...)
			case pos == len(res.GetResults())-1:
				res.Results = append(res.GetResults(), dist)
			case pos >= 0:
				res.Results = append(res.GetResults()[:pos+1], res.GetResults()[pos:]...)
				res.Results[pos+1] = dist
			}
			if len(res.GetResults()) > num && num != 0 {
				res.Results = res.GetResults()[:num]
			}
		}
	}
}

func (s *server) StreamSearch(stream vald.Vald_StreamSearchServer) error {
	return grpc.BidirectionalStream(stream, s.streamConcurrency,
		func() interface{} { return new(payload.Search_Request) },
		func(ctx context.Context, data interface{}) (interface{}, error) {
			return s.Search(ctx, data.(*payload.Search_Request))
		})
}

func (s *server) StreamSearchByID(stream vald.Vald_StreamSearchByIDServer) error {
	return grpc.BidirectionalStream(stream, s.streamConcurrency,
		func() interface{} { return new(payload.Search_IDRequest) },
		func(ctx context.Context, data interface{}) (interface{}, error) {
			return s.SearchByID(ctx, data.(*payload.Search_IDRequest))
		})
}

func (s *server) Insert(ctx context.Context, vec *payload.Object_Vector) (ce *payload.Empty, err error) {
	meta := vec.GetId()
	uuid, err := s.metadata.GetUUID(ctx, meta)
	if err == nil || len(uuid) != 0 {
		return nil, status.WrapWithAlreadyExists(fmt.Sprintf("Insert API meta %s already exists", meta),
			errors.ErrMetaDataAlreadyExists(meta, uuid), info.Get())
	}

	uuid = fuid.String()
	err = s.metadata.SetUUIDandMeta(ctx, uuid, meta)
	if err != nil {
		log.Error(err)
		return nil, status.WrapWithInternal(fmt.Sprintf("Insert API meta %s & uuid %s couldn't store", meta, uuid), err, info.Get())
	}
	vec.Id = uuid
	mu := new(sync.Mutex)
	targets := make([]string, 0, s.replica)
	err = s.gateway.DoMulti(ctx, s.replica, func(ctx context.Context, target string, ac agent.AgentClient, copts ...grpc.CallOption) (err error) {
		_, err = ac.Insert(ctx, vec, copts...)
		if err != nil {
			return err
		}
		target = strings.SplitN(target, ":", 2)[0]
		mu.Lock()
		targets = append(targets, target)
		mu.Unlock()
		return nil
	})
	if err != nil {
		log.Error(err)
		return nil, status.WrapWithInternal(fmt.Sprintf("Insert API failed to Execute DoMulti error = %s", err.Error()), err, info.Get())
	}
	if s.backup != nil {
		vecs := &payload.Backup_MetaVector{
			Uuid:   uuid,
			Meta:   meta,
			Vector: vec.GetVector(),
			Ips:    targets,
		}
		err = s.backup.Register(ctx, vecs)
		if err != nil {
			log.Error(err)
			return nil, status.WrapWithInternal(fmt.Sprintf("Insert API failed to Backup Vectors %#v", vecs), err, info.Get())
		}
	}
	return new(payload.Empty), nil
}

func (s *server) StreamInsert(stream vald.Vald_StreamInsertServer) error {
	return grpc.BidirectionalStream(stream, s.streamConcurrency,
		func() interface{} { return new(payload.Object_Vector) },
		func(ctx context.Context, data interface{}) (interface{}, error) {
			return s.Insert(ctx, data.(*payload.Object_Vector))
		})
}

func (s *server) MultiInsert(ctx context.Context, vecs *payload.Object_Vectors) (res *payload.Empty, err error) {
	metaMap := make(map[string]string)
	metas := make([]string, 0, len(vecs.GetVectors()))
	for i, vec := range vecs.GetVectors() {
		uuid := fuid.String()
		meta := vec.GetId()
		metaMap[uuid] = meta
		metas = append(metas, meta)
		vecs.Vectors[i].Id = uuid
	}
	uuids, err := s.metadata.GetMetas(ctx, metas...)
	if err == nil {
		for i, meta := range metas {
			if len(uuids) > i && len(uuids[i]) != 0 {
				err = errors.Wrap(err, errors.ErrMetaDataAlreadyExists(meta, uuids[i]).Error())
			}
		}
		if err != nil {
			return nil, status.WrapWithAlreadyExists(fmt.Sprintf("MultiInsert API failed metadata already exists uuids metas = %v", metas), err, info.Get())
		}
	}

	mu := new(sync.Mutex)
	targets := make([]string, 0, s.replica)
	gerr := s.gateway.DoMulti(ctx, s.replica, func(ctx context.Context, target string, ac agent.AgentClient, copts ...grpc.CallOption) (err error) {
		_, err = ac.MultiInsert(ctx, vecs, copts...)
		if err != nil {
			return err
		}
		target = strings.SplitN(target, ":", 2)[0]
		mu.Lock()
		targets = append(targets, target)
		mu.Unlock()
		return nil
	})
	if gerr != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("MultiInsert API failed request %#v", vecs), errors.Wrap(gerr, err.Error()), info.Get())
	}

	err = s.metadata.SetUUIDandMetas(ctx, metaMap)
	if err != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("MultiInsert API failed SetUUIDandMetas %#v", metaMap), err, info.Get())
	}

	if s.backup != nil {
		mvecs := new(payload.Backup_MetaVectors)
		mvecs.Vectors = make([]*payload.Backup_MetaVector, 0, len(vecs.GetVectors()))
		for _, vec := range vecs.GetVectors() {
			uuid := vec.GetId()
			mvecs.Vectors = append(mvecs.Vectors, &payload.Backup_MetaVector{
				Uuid:   uuid,
				Meta:   metaMap[uuid],
				Vector: vec.GetVector(),
				Ips:    targets,
			})
		}
		err = s.backup.RegisterMultiple(ctx, mvecs)
		if err != nil {
			return nil, status.WrapWithInternal(fmt.Sprintf("MultiInsert API failed RegisterMultiple %#v", mvecs), err, info.Get())
		}
	}
	return new(payload.Empty), nil
}

func (s *server) Update(ctx context.Context, vec *payload.Object_Vector) (res *payload.Empty, err error) {
	meta := vec.GetId()
	uuid, err := s.metadata.GetUUID(ctx, meta)
	if err != nil {
		return nil, status.WrapWithNotFound(fmt.Sprintf("Update API failed GetUUID meta = %s", meta), err, info.Get())
	}
	vec.Id = uuid
	locs, err := s.backup.GetLocation(ctx, uuid)
	if err != nil {
		return nil, status.WrapWithNotFound(fmt.Sprintf("Update API failed GetLocation meta = %s, uuid = %s", meta, uuid), err, info.Get())
	}
	lmap := make(map[string]struct{}, len(locs))
	for _, loc := range locs {
		lmap[loc] = struct{}{}
	}
	err = s.gateway.BroadCast(ctx, func(ctx context.Context, target string, ac agent.AgentClient, copts ...grpc.CallOption) error {
		target = strings.SplitN(target, ":", 2)[0]
		_, ok := lmap[target]
		if ok {
			_, err = ac.Update(ctx, vec, copts...)
			if err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("Update API failed request %#v", vec), err, info.Get())
	}
	mvec := &payload.Backup_MetaVector{
		Uuid:   uuid,
		Meta:   meta,
		Vector: vec.GetVector(),
		Ips:    locs,
	}
	err = s.backup.Register(ctx, mvec)
	if err != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("Update API failed backup %#v", vec), err, info.Get())
	}

	return new(payload.Empty), nil
}

func (s *server) StreamUpdate(stream vald.Vald_StreamUpdateServer) error {
	return grpc.BidirectionalStream(stream, s.streamConcurrency,
		func() interface{} { return new(payload.Object_Vector) },
		func(ctx context.Context, data interface{}) (interface{}, error) {
			return s.Update(ctx, data.(*payload.Object_Vector))
		})
}

func (s *server) MultiUpdate(ctx context.Context, vecs *payload.Object_Vectors) (res *payload.Empty, err error) {
	ids := make([]string, 0, len(vecs.GetVectors()))
	for _, vec := range vecs.GetVectors() {
		ids = append(ids, vec.GetId())
	}
	_, err = s.MultiRemove(ctx, &payload.Object_IDs{
		Ids: ids,
	})
	if err != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("MultiUpdate API failed Remove request %#v", ids), err, info.Get())
	}
	_, err = s.MultiInsert(ctx, vecs)
	if err != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("MultiUpdate API failed Insert request %#v", vecs), err, info.Get())
	}
	return new(payload.Empty), nil
}

func (s *server) Upsert(ctx context.Context, vec *payload.Object_Vector) (*payload.Empty, error) {
	meta := vec.GetId()
	uuid, err := s.metadata.GetUUID(ctx, meta)
	if err != nil || len(uuid) == 0 {
		return s.Insert(ctx, vec)
	}
	return s.Update(ctx, vec)
}

func (s *server) StreamUpsert(stream vald.Vald_StreamUpsertServer) error {
	return grpc.BidirectionalStream(stream, s.streamConcurrency,
		func() interface{} { return new(payload.Object_Vector) },
		func(ctx context.Context, data interface{}) (interface{}, error) {
			return s.Upsert(ctx, data.(*payload.Object_Vector))
		})
}

func (s *server) MultiUpsert(ctx context.Context, vecs *payload.Object_Vectors) (*payload.Empty, error) {
	metas := make([]string, 0, len(vecs.GetVectors()))
	for _, vec := range vecs.GetVectors() {
		metas = append(metas, vec.GetId())
	}
	uuids, err := s.metadata.GetUUIDs(ctx, metas...)
	if err == nil {
		for i := range uuids {
			if len(uuids) > i && len(uuids[i]) != 0 {
				return s.MultiUpdate(ctx, vecs)
			}
		}
	}
	return s.MultiInsert(ctx, vecs)
}

func (s *server) Remove(ctx context.Context, id *payload.Object_ID) (*payload.Empty, error) {
	meta := id.GetId()
	uuid, err := s.metadata.GetUUID(ctx, meta)
	if err != nil {
		return nil, status.WrapWithNotFound(fmt.Sprintf("Remove API meta %s's uuid not found", meta), err, info.Get())
	}
	locs, err := s.backup.GetLocation(ctx, uuid)
	if err != nil {
		return nil, status.WrapWithNotFound(fmt.Sprintf("Remove API failed GetLocation meta = %s, uuid = %s", meta, uuid), err, info.Get())
	}
	lmap := make(map[string]struct{}, len(locs))
	for _, loc := range locs {
		lmap[loc] = struct{}{}
	}
	err = s.gateway.BroadCast(ctx, func(ctx context.Context, target string, ac agent.AgentClient, copts ...grpc.CallOption) error {
		_, ok := lmap[target]
		if ok {
			_, err = ac.Remove(ctx, &payload.Object_ID{
				Id: uuid,
			}, copts...)
			if err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("Remove API failed request uuid %s", uuid), err, info.Get())
	}
	_, err = s.metadata.DeleteMeta(ctx, uuid)
	if err != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("Remove API failed Delete metadata uuid = %s", uuid), err, info.Get())
	}
	err = s.backup.Remove(ctx, uuid)
	if err != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("Remove API failed to Remove backup uuid = %s", uuid), err, info.Get())
	}
	return new(payload.Empty), nil
}

func (s *server) StreamRemove(stream vald.Vald_StreamRemoveServer) error {
	return grpc.BidirectionalStream(stream, s.streamConcurrency,
		func() interface{} { return new(payload.Object_ID) },
		func(ctx context.Context, data interface{}) (interface{}, error) {
			return s.Remove(ctx, data.(*payload.Object_ID))
		})
}

func (s *server) MultiRemove(ctx context.Context, ids *payload.Object_IDs) (res *payload.Empty, err error) {
	uuids, err := s.metadata.GetUUIDs(ctx, ids.GetIds()...)
	if err != nil {
		return nil, status.WrapWithNotFound(fmt.Sprintf("MultiRemove API meta datas %v's uuid not found", ids.GetIds()), err, info.Get())
	}
	lmap := make(map[string][]string, s.gateway.GetAgentCount())
	for _, uuid := range uuids {
		locs, err := s.backup.GetLocation(ctx, uuid)
		if err != nil {
			return nil, status.WrapWithNotFound(fmt.Sprintf("MultiRemove API failed to get uuid %s's location ", uuid), err, info.Get())
		}
		for _, loc := range locs {
			lmap[loc] = append(lmap[loc], uuid)
		}
	}
	err = s.gateway.BroadCast(ctx, func(ctx context.Context, target string, ac agent.AgentClient, copts ...grpc.CallOption) error {
		uuids, ok := lmap[target]
		if ok {
			_, err := ac.MultiRemove(ctx, &payload.Object_IDs{
				Ids: uuids,
			}, copts...)
			if err != nil {
				return err
			}
		}
		return nil
	})
	if err != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("MultiRemove API failed to request uuids %v metas %v ", uuids, ids.GetIds()), err, info.Get())
	}
	_, err = s.metadata.DeleteMetas(ctx, uuids...)
	if err != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("MultiRemove API failed to DeleteMetas uuids %v ", uuids), err, info.Get())
	}
	err = s.backup.RemoveMultiple(ctx, uuids...)
	if err != nil {
		return nil, status.WrapWithInternal(fmt.Sprintf("MultiRemove API failed to Remove backup uuids %v ", uuids), err, info.Get())
	}
	return new(payload.Empty), nil
}

func (s *server) GetObject(ctx context.Context, id *payload.Object_ID) (vec *payload.Backup_MetaVector, err error) {
	meta := id.GetId()
	uuid, err := s.metadata.GetUUID(ctx, meta)
	if err != nil {
		return nil, status.WrapWithNotFound(fmt.Sprintf("GetObject API meta %s's uuid not found", meta), err, info.Get())
	}
	vec, err = s.backup.GetObject(ctx, uuid)
	if err != nil {
		return nil, status.WrapWithNotFound(fmt.Sprintf("GetObject API meta %s uuid %s Object not found", meta, uuid), err, info.Get())
	}
	return vec, nil
}

func (s *server) StreamGetObject(stream vald.Vald_StreamGetObjectServer) error {
	return grpc.BidirectionalStream(stream, s.streamConcurrency,
		func() interface{} { return new(payload.Object_ID) },
		func(ctx context.Context, data interface{}) (interface{}, error) {
			return s.GetObject(ctx, data.(*payload.Object_ID))
		})
}