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.

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