goc/pkg/cover/store.go

/*
 Copyright 2020 Qiniu Cloud (qiniu.com)

 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 cover

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	log "github.com/sirupsen/logrus"
)

var ErrServiceAlreadyRegistered = errors.New("service already registered")

// Store persistents the registered service information
type Store interface {
	// Add adds the given service to store
	Add(s ServiceUnderTest) error

	// Get returns the registered service information with the given service's name
	Get(name string) []string

	// Get returns all the registered service information as a map
	GetAll() map[string][]string

	// Init cleanup all the registered service information
	Init() error

	// Set stores the services information into internal state
	Set(services map[string][]string) error

	// Remove the service from the store by address
	Remove(addr string) error
}

// fileStore holds the registered services into memory and persistent to a local file
type fileStore struct {
	mu             sync.RWMutex
	persistentFile string

	memoryStore Store
}

// NewFileStore creates a store using local file
func NewFileStore(persistenceFile string) (store Store, err error) {
	path, err := filepath.Abs(persistenceFile)
	if err != nil {
		return nil, err
	}
	err = os.MkdirAll(filepath.Dir(path), os.ModePerm)
	if err != nil {
		return nil, err
	}
	l := &fileStore{
		persistentFile: path,
		memoryStore:    NewMemoryStore(),
	}

	if err := l.load(); err != nil {
		log.Fatalf("load failed, file: %s, err: %v", l.persistentFile, err)
	}

	return l, nil
}

// Add adds the given service to file Store
func (l *fileStore) Add(s ServiceUnderTest) error {
	if err := l.memoryStore.Add(s); err != nil {
		return err
	}

	// persistent to local store
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.appendToFile(s)
}

// Get returns the registered service information with the given name
func (l *fileStore) Get(name string) []string {
	return l.memoryStore.Get(name)
}

// Get returns all the registered service information
func (l *fileStore) GetAll() map[string][]string {
	return l.memoryStore.GetAll()
}

// Remove the service from the memory store and the file store
func (l *fileStore) Remove(addr string) error {
	err := l.memoryStore.Remove(addr)
	if err != nil {
		return err
	}

	return l.Set(l.memoryStore.GetAll())
}

// Init cleanup all the registered service information
// and the local persistent file
func (l *fileStore) Init() error {
	if err := l.memoryStore.Init(); err != nil {
		return err
	}

	l.mu.Lock()
	defer l.mu.Unlock()
	if err := os.Remove(l.persistentFile); err != nil && !os.IsNotExist(err) {
		return fmt.Errorf("failed to delete file %s, err: %v", l.persistentFile, err)
	}

	return nil
}

// load all registered service from file to memory
func (l *fileStore) load() error {
	var svrsMap = make(map[string][]string, 0)

	f, err := os.Open(l.persistentFile)
	if err != nil {
		if os.IsNotExist(err) {
			return nil
		}
		return fmt.Errorf("failed to open file, path: %s, err: %v", l.persistentFile, err)
	}
	defer f.Close()

	ns := bufio.NewScanner(f)
	for ns.Scan() {
		line := ns.Text()
		ss := strings.FieldsFunc(line, split)

		// TODO: use regex
		if len(ss) == 2 {
			if urls, ok := svrsMap[ss[0]]; ok {
				urls = append(urls, ss[1])
				svrsMap[ss[0]] = urls
			} else {
				svrsMap[ss[0]] = []string{ss[1]}
			}
		}
	}

	if err := ns.Err(); err != nil {
		return fmt.Errorf("read file failed, file: %s, err: %v", l.persistentFile, err)
	}

	// set information to memory
	l.memoryStore.Set(svrsMap)
	return nil
}

func (l *fileStore) Set(services map[string][]string) error {
	l.mu.Lock()
	defer l.mu.Unlock()

	// no error will return from memorystore.set
	err := l.memoryStore.Set(services)
	if err != nil {
		return err
	}

	f, err := os.OpenFile(l.persistentFile, os.O_TRUNC|os.O_WRONLY|os.O_CREATE, 0600)
	if err != nil {
		return err
	}

	s := ""
	for name, addrs := range services {
		for _, addr := range addrs {
			s += fmt.Sprintf("%s&%s\n", name, addr)
		}
	}

	_, err = f.WriteString(s)
	if err != nil {
		return err
	}

	return f.Sync()
}

func (l *fileStore) appendToFile(s ServiceUnderTest) error {
	f, err := os.OpenFile(l.persistentFile, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0600)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = f.WriteString(format(s) + "\n")
	if err != nil {
		return err
	}

	f.Sync()
	return nil
}

func format(s ServiceUnderTest) string {
	return fmt.Sprintf("%s&%s", s.Name, s.Address)
}

func split(r rune) bool {
	return r == '&'
}

// memoryStore holds the registered services only into memory
type memoryStore struct {
	mu          sync.RWMutex
	servicesMap map[string][]string
}

// NewMemoryStore creates a memory store
func NewMemoryStore() Store {
	return &memoryStore{
		servicesMap: make(map[string][]string, 0),
	}
}

// Add adds the given service to MemoryStore
func (l *memoryStore) Add(s ServiceUnderTest) error {
	l.mu.Lock()
	defer l.mu.Unlock()
	// load to memory
	if addrs, ok := l.servicesMap[s.Name]; ok {
		for _, addr := range addrs {
			if addr == s.Address {
				log.Printf("service registered already, name: %s, address: %s", s.Name, s.Address)
				return ErrServiceAlreadyRegistered
			}
		}
		addrs = append(addrs, s.Address)
		l.servicesMap[s.Name] = addrs
	} else {
		l.servicesMap[s.Name] = []string{s.Address}
	}

	return nil
}

// Get returns the registered service information with the given name
func (l *memoryStore) Get(name string) []string {
	l.mu.RLock()
	defer l.mu.RUnlock()
	return l.servicesMap[name]
}

// Get returns all the registered service information
func (l *memoryStore) GetAll() map[string][]string {
	res := make(map[string][]string)
	l.mu.RLock()
	defer l.mu.RUnlock()
	for k,v := range l.servicesMap {
		res[k] = append(make([]string,0,len(v)), v...)
	}
	return res
}

// Init cleanup all the registered service information
// and the local persistent file
func (l *memoryStore) Init() error {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.servicesMap = make(map[string][]string, 0)
	return nil
}

func (l *memoryStore) Set(services map[string][]string) error {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.servicesMap = services

	return nil
}

// Remove one service from the memory store
// if service is not fount, return "no service found" error
func (l *memoryStore) Remove(removeAddr string) error {
	l.mu.Lock()
	defer l.mu.Unlock()

	flag := false
	for name, addrs := range l.servicesMap {
		newAddrs := make([]string, 0)
		for _, addr := range addrs {
			if removeAddr != addr {
				newAddrs = append(newAddrs, addr)
			} else {
				flag = true
			}
		}
		// if no services left, remove by name
		if len(newAddrs) == 0 {
			delete(l.servicesMap, name)
		} else {
			l.servicesMap[name] = newAddrs
		}
	}

	if !flag {
		return fmt.Errorf("no service found")
	}

	return nil
}