godocdb/storage/storage.go

package storage

import (
	"fmt"
	"io/ioutil"
	"os"
	"sync"
)

type DbPolicy int

const (
	DbPolicy_Tidy  DbPolicy = iota // all store in one db
	DbPolicy_Loosy                 // store in different db
)

const (
	DefaultDbName = "default"
)

type StorageManager struct {
	policy    DbPolicy
	dbrootDir string
	db        map[string]*LevelDBStorage
}

var (
	instance *StorageManager
	once     sync.Once
)

// 修改工厂函数实现单例模式
func NewStorageManager(dbrootDir string, dbPolicy int) (*StorageManager, error) {
	var initErr error
	once.Do(func() {
		instance = &StorageManager{
			policy:    DbPolicy(dbPolicy),
			dbrootDir: dbrootDir,
			db:        make(map[string]*LevelDBStorage),
		}

		// 预初始化默认数据库
		if _, err := instance.GetStorage(DefaultDbName); err != nil {
			initErr = err
		}
	})

	if initErr != nil {
		return nil, initErr
	}
	return instance, nil
}

// 新增获取单例实例的方法
func GetInstance() *StorageManager {
	return instance
}

func (sm *StorageManager) GetStorage(dbName string) (*LevelDBStorage, error) {
	if dbName == "" {
		return nil, fmt.Errorf("db name is empty")
	}
	if sm.policy == DbPolicy_Tidy {
		dbName = DefaultDbName
	}
	_, exists := sm.db[dbName]
	if !exists {
		db, err := NewLevelDBStorage(sm.dbrootDir + "/" + dbName)
		if err != nil {
			return nil, err
		}
		sm.db[dbName] = db
	}
	return sm.db[dbName], nil
}

func (sm *StorageManager) isLevelDB(dirp string) bool {
	if _, err := os.Stat(dirp); os.IsNotExist(err) {
		return false
	}
	files := []string{"CURRENT", "LOCK", "LOG", "LOG.old", "MANIFEST-"}
	for _, file := range files {
		if _, err := os.Stat(dirp + "/" + file); err == nil {
			return true
		}
	}
	return false
}

func (sm *StorageManager) Open() {
	// 得到dbrootDir下的所有子目录
	dirs, err := ioutil.ReadDir(sm.dbrootDir)
	if err != nil {
		panic(err)
	}
	for _, dir := range dirs {
		if dir.IsDir() {
			dbPath := sm.dbrootDir + "/" + dir.Name()
			if sm.isLevelDB(dbPath) {
				db, err := NewLevelDBStorage(dbPath)
				if err != nil {
					panic(err)
				}
				sm.db[dir.Name()] = db
			}

		}
	}
}
func (sm *StorageManager) Close() {
	for _, db := range sm.db {
		db.Close()
	}
}