You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

396 lines
8.9 KiB

3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
3 years ago
  1. //go:build tikv
  2. // +build tikv
  3. package tikv
  4. import (
  5. "bytes"
  6. "context"
  7. "crypto/sha1"
  8. "fmt"
  9. "io"
  10. "strings"
  11. "github.com/chrislusf/seaweedfs/weed/filer"
  12. "github.com/chrislusf/seaweedfs/weed/glog"
  13. "github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
  14. "github.com/chrislusf/seaweedfs/weed/util"
  15. "github.com/tikv/client-go/v2/txnkv"
  16. )
  17. var (
  18. _ filer.FilerStore = ((*TikvStore)(nil))
  19. )
  20. func init() {
  21. filer.Stores = append(filer.Stores, &TikvStore{})
  22. }
  23. type TikvStore struct {
  24. client *txnkv.Client
  25. deleteRangeConcurrency int
  26. onePC bool
  27. }
  28. // Basic APIs
  29. func (store *TikvStore) GetName() string {
  30. return "tikv"
  31. }
  32. func (store *TikvStore) Initialize(config util.Configuration, prefix string) error {
  33. pdAddrs := []string{}
  34. pdAddrsStr := config.GetString(prefix + "pdaddrs")
  35. for _, item := range strings.Split(pdAddrsStr, ",") {
  36. pdAddrs = append(pdAddrs, strings.TrimSpace(item))
  37. }
  38. drc := config.GetInt(prefix + "deleterange_concurrency")
  39. if drc <= 0 {
  40. drc = 1
  41. }
  42. store.onePC = config.GetBool(prefix + "enable_1pc")
  43. store.deleteRangeConcurrency = drc
  44. return store.initialize(pdAddrs)
  45. }
  46. func (store *TikvStore) initialize(pdAddrs []string) error {
  47. client, err := txnkv.NewClient(pdAddrs)
  48. store.client = client
  49. return err
  50. }
  51. func (store *TikvStore) Shutdown() {
  52. err := store.client.Close()
  53. if err != nil {
  54. glog.V(0).Infof("Shutdown TiKV client got error: %v", err)
  55. }
  56. }
  57. // ~ Basic APIs
  58. // Entry APIs
  59. func (store *TikvStore) InsertEntry(ctx context.Context, entry *filer.Entry) error {
  60. dir, name := entry.DirAndName()
  61. key := generateKey(dir, name)
  62. value, err := entry.EncodeAttributesAndChunks()
  63. if err != nil {
  64. return fmt.Errorf("encoding %s %+v: %v", entry.FullPath, entry.Attr, err)
  65. }
  66. txn, err := store.getTxn(ctx)
  67. if err != nil {
  68. return err
  69. }
  70. err = txn.RunInTxn(func(txn *txnkv.KVTxn) error {
  71. return txn.Set(key, value)
  72. })
  73. if err != nil {
  74. return fmt.Errorf("persisting %s : %v", entry.FullPath, err)
  75. }
  76. return nil
  77. }
  78. func (store *TikvStore) UpdateEntry(ctx context.Context, entry *filer.Entry) error {
  79. return store.InsertEntry(ctx, entry)
  80. }
  81. func (store *TikvStore) FindEntry(ctx context.Context, path util.FullPath) (*filer.Entry, error) {
  82. dir, name := path.DirAndName()
  83. key := generateKey(dir, name)
  84. txn, err := store.getTxn(ctx)
  85. if err != nil {
  86. return nil, err
  87. }
  88. var value []byte = nil
  89. err = txn.RunInTxn(func(txn *txnkv.KVTxn) error {
  90. val, err := txn.Get(context.TODO(), key)
  91. if err == nil {
  92. value = val
  93. }
  94. return err
  95. })
  96. if isNotExists(err) || value == nil {
  97. return nil, filer_pb.ErrNotFound
  98. }
  99. if err != nil {
  100. return nil, fmt.Errorf("get %s : %v", path, err)
  101. }
  102. entry := &filer.Entry{
  103. FullPath: path,
  104. }
  105. err = entry.DecodeAttributesAndChunks(value)
  106. if err != nil {
  107. return entry, fmt.Errorf("decode %s : %v", entry.FullPath, err)
  108. }
  109. return entry, nil
  110. }
  111. func (store *TikvStore) DeleteEntry(ctx context.Context, path util.FullPath) error {
  112. dir, name := path.DirAndName()
  113. key := generateKey(dir, name)
  114. txn, err := store.getTxn(ctx)
  115. if err != nil {
  116. return err
  117. }
  118. err = txn.RunInTxn(func(txn *txnkv.KVTxn) error {
  119. return txn.Delete(key)
  120. })
  121. if err != nil {
  122. return fmt.Errorf("delete %s : %v", path, err)
  123. }
  124. return nil
  125. }
  126. // ~ Entry APIs
  127. // Directory APIs
  128. func (store *TikvStore) DeleteFolderChildren(ctx context.Context, path util.FullPath) error {
  129. directoryPrefix := genDirectoryKeyPrefix(path, "")
  130. txn, err := store.getTxn(ctx)
  131. if err != nil {
  132. return err
  133. }
  134. var (
  135. startKey []byte = nil
  136. endKey []byte = nil
  137. )
  138. err = txn.RunInTxn(func(txn *txnkv.KVTxn) error {
  139. iter, err := txn.Iter(directoryPrefix, nil)
  140. if err != nil {
  141. return err
  142. }
  143. defer iter.Close()
  144. for iter.Valid() {
  145. key := iter.Key()
  146. endKey = key
  147. if !bytes.HasPrefix(key, directoryPrefix) {
  148. break
  149. }
  150. if startKey == nil {
  151. startKey = key
  152. }
  153. err = iter.Next()
  154. if err != nil {
  155. return err
  156. }
  157. }
  158. // Only one Key matched just delete it.
  159. if startKey != nil && bytes.Equal(startKey, endKey) {
  160. return txn.Delete(startKey)
  161. }
  162. return nil
  163. })
  164. if err != nil {
  165. return fmt.Errorf("delete %s : %v", path, err)
  166. }
  167. if startKey != nil && endKey != nil && !bytes.Equal(startKey, endKey) {
  168. // has startKey and endKey and they are not equals, so use delete range
  169. _, err = store.client.DeleteRange(context.Background(), startKey, endKey, store.deleteRangeConcurrency)
  170. if err != nil {
  171. return fmt.Errorf("delete %s : %v", path, err)
  172. }
  173. }
  174. return err
  175. }
  176. func (store *TikvStore) ListDirectoryEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, eachEntryFunc filer.ListEachEntryFunc) (string, error) {
  177. return store.ListDirectoryPrefixedEntries(ctx, dirPath, startFileName, includeStartFile, limit, "", eachEntryFunc)
  178. }
  179. func (store *TikvStore) ListDirectoryPrefixedEntries(ctx context.Context, dirPath util.FullPath, startFileName string, includeStartFile bool, limit int64, prefix string, eachEntryFunc filer.ListEachEntryFunc) (string, error) {
  180. lastFileName := ""
  181. directoryPrefix := genDirectoryKeyPrefix(dirPath, prefix)
  182. lastFileStart := directoryPrefix
  183. if startFileName != "" {
  184. lastFileStart = genDirectoryKeyPrefix(dirPath, startFileName)
  185. }
  186. txn, err := store.getTxn(ctx)
  187. if err != nil {
  188. return lastFileName, err
  189. }
  190. err = txn.RunInTxn(func(txn *txnkv.KVTxn) error {
  191. iter, err := txn.Iter(lastFileStart, nil)
  192. if err != nil {
  193. return err
  194. }
  195. defer iter.Close()
  196. i := int64(0)
  197. first := true
  198. for iter.Valid() {
  199. if first {
  200. first = false
  201. if !includeStartFile {
  202. if iter.Valid() {
  203. // Check first item is lastFileStart
  204. if bytes.Equal(iter.Key(), lastFileStart) {
  205. // Is lastFileStart and not include start file, just
  206. // ignore it.
  207. err = iter.Next()
  208. if err != nil {
  209. return err
  210. }
  211. continue
  212. }
  213. }
  214. }
  215. }
  216. // Check for limitation
  217. if limit > 0 {
  218. i++
  219. if i > limit {
  220. break
  221. }
  222. }
  223. // Validate key prefix
  224. key := iter.Key()
  225. if !bytes.HasPrefix(key, directoryPrefix) {
  226. break
  227. }
  228. value := iter.Value()
  229. // Start process
  230. fileName := getNameFromKey(key)
  231. if fileName != "" {
  232. // Got file name, then generate the Entry
  233. entry := &filer.Entry{
  234. FullPath: util.NewFullPath(string(dirPath), fileName),
  235. }
  236. // Update lastFileName
  237. lastFileName = fileName
  238. // Check for decode value.
  239. if decodeErr := entry.DecodeAttributesAndChunks(value); decodeErr != nil {
  240. // Got error just return the error
  241. glog.V(0).Infof("list %s : %v", entry.FullPath, err)
  242. return err
  243. }
  244. // Run for each callback if return false just break the iteration
  245. if !eachEntryFunc(entry) {
  246. break
  247. }
  248. }
  249. // End process
  250. err = iter.Next()
  251. if err != nil {
  252. return err
  253. }
  254. }
  255. return nil
  256. })
  257. if err != nil {
  258. return lastFileName, fmt.Errorf("prefix list %s : %v", dirPath, err)
  259. }
  260. return lastFileName, nil
  261. }
  262. // ~ Directory APIs
  263. // Transaction Related APIs
  264. func (store *TikvStore) BeginTransaction(ctx context.Context) (context.Context, error) {
  265. tx, err := store.client.Begin()
  266. if err != nil {
  267. return ctx, err
  268. }
  269. if store.onePC {
  270. tx.SetEnable1PC(store.onePC)
  271. }
  272. return context.WithValue(ctx, "tx", tx), nil
  273. }
  274. func (store *TikvStore) CommitTransaction(ctx context.Context) error {
  275. if tx, ok := ctx.Value("tx").(*txnkv.KVTxn); ok {
  276. return tx.Commit(context.Background())
  277. }
  278. return nil
  279. }
  280. func (store *TikvStore) RollbackTransaction(ctx context.Context) error {
  281. if tx, ok := ctx.Value("tx").(*txnkv.KVTxn); ok {
  282. return tx.Rollback()
  283. }
  284. return nil
  285. }
  286. // ~ Transaction Related APIs
  287. // Transaction Wrapper
  288. type TxnWrapper struct {
  289. *txnkv.KVTxn
  290. inContext bool
  291. }
  292. func (w *TxnWrapper) RunInTxn(f func(txn *txnkv.KVTxn) error) error {
  293. err := f(w.KVTxn)
  294. if !w.inContext {
  295. if err != nil {
  296. w.KVTxn.Rollback()
  297. return err
  298. }
  299. w.KVTxn.Commit(context.Background())
  300. return nil
  301. }
  302. return err
  303. }
  304. func (store *TikvStore) getTxn(ctx context.Context) (*TxnWrapper, error) {
  305. if tx, ok := ctx.Value("tx").(*txnkv.KVTxn); ok {
  306. return &TxnWrapper{tx, true}, nil
  307. }
  308. txn, err := store.client.Begin()
  309. if err != nil {
  310. return nil, err
  311. }
  312. if store.onePC {
  313. txn.SetEnable1PC(store.onePC)
  314. }
  315. return &TxnWrapper{txn, false}, nil
  316. }
  317. // ~ Transaction Wrapper
  318. // Encoding Functions
  319. func hashToBytes(dir string) []byte {
  320. h := sha1.New()
  321. io.WriteString(h, dir)
  322. b := h.Sum(nil)
  323. return b
  324. }
  325. func generateKey(dirPath, fileName string) []byte {
  326. key := hashToBytes(dirPath)
  327. key = append(key, []byte(fileName)...)
  328. return key
  329. }
  330. func getNameFromKey(key []byte) string {
  331. return string(key[sha1.Size:])
  332. }
  333. func genDirectoryKeyPrefix(fullpath util.FullPath, startFileName string) (keyPrefix []byte) {
  334. keyPrefix = hashToBytes(string(fullpath))
  335. if len(startFileName) > 0 {
  336. keyPrefix = append(keyPrefix, []byte(startFileName)...)
  337. }
  338. return keyPrefix
  339. }
  340. func isNotExists(err error) bool {
  341. if err == nil {
  342. return false
  343. }
  344. if err.Error() == "not exist" {
  345. return true
  346. }
  347. return false
  348. }
  349. // ~ Encoding Functions