kraken

1
// Copyright (c) 2016-2019 Uber Technologies, Inc.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14
package agentstorage
15

16
import (
17
	"fmt"
18
	"os"
19

20
	"github.com/uber-go/tally"
21
	"github.com/willf/bitset"
22

23
	"github.com/uber/kraken/core"
24
	"github.com/uber/kraken/lib/store"
25
	"github.com/uber/kraken/lib/store/metadata"
26
	"github.com/uber/kraken/lib/torrent/storage"
27
	"github.com/uber/kraken/tracker/metainfoclient"
28
)
29

30
// TorrentArchive is capable of initializing torrents in the download directory
31
// and serving torrents from either the download or cache directory.
32
type TorrentArchive struct {
33
	stats          tally.Scope
34
	cads           *store.CADownloadStore
35
	metaInfoClient metainfoclient.Client
36
}
37

38
// NewTorrentArchive creates a new TorrentArchive.
39
func NewTorrentArchive(
40
	stats tally.Scope,
41
	cads *store.CADownloadStore,
42
	mic metainfoclient.Client) *TorrentArchive {
43

44
	stats = stats.Tagged(map[string]string{
45
		"module": "agenttorrentarchive",
46
	})
47

48
	return &TorrentArchive{stats, cads, mic}
49
}
50

51
// Stat returns TorrentInfo for the given digest. Returns os.ErrNotExist if the
52
// file does not exist. Ignores namespace.
53
func (a *TorrentArchive) Stat(namespace string, d core.Digest) (*storage.TorrentInfo, error) {
54
	var tm metadata.TorrentMeta
55
	if err := a.cads.Any().GetMetadata(d.Hex(), &tm); err != nil {
56
		return nil, err
57
	}
58
	var psm pieceStatusMetadata
59
	if err := a.cads.Any().GetMetadata(d.Hex(), &psm); err != nil {
60
		return nil, err
61
	}
62
	b := bitset.New(uint(len(psm.pieces)))
63
	for i, p := range psm.pieces {
64
		if p.status == _complete {
65
			b.Set(uint(i))
66
		}
67
	}
68
	return storage.NewTorrentInfo(tm.MetaInfo, b), nil
69
}
70

71
// CreateTorrent returns a Torrent for either an existing metainfo / file on
72
// disk, or downloads metainfo and initializes the file. Returns ErrNotFound
73
// if no metainfo was found.
74
func (a *TorrentArchive) CreateTorrent(namespace string, d core.Digest) (storage.Torrent, error) {
75
	var tm metadata.TorrentMeta
76
	if err := a.cads.Any().GetMetadata(d.Hex(), &tm); os.IsNotExist(err) {
77
		downloadTimer := a.stats.Timer("metainfo_download").Start()
78
		mi, err := a.metaInfoClient.Download(namespace, d)
79
		if err != nil {
80
			if err == metainfoclient.ErrNotFound {
81
				return nil, storage.ErrNotFound
82
			}
83
			return nil, fmt.Errorf("download metainfo: %s", err)
84
		}
85
		downloadTimer.Stop()
86

87
		// There's a race condition here, but it's "okay"... Basically, we could
88
		// initialize a download file with metainfo that is rejected by file store,
89
		// because someone else beats us to it. However, we catch a lucky break
90
		// because the only piece of metainfo we use is file length -- which digest
91
		// is derived from, so it's "okay".
92
		createErr := a.cads.CreateDownloadFile(mi.Digest().Hex(), mi.Length())
93
		if createErr != nil &&
94
			!(a.cads.InDownloadError(createErr) || a.cads.InCacheError(createErr)) {
95
			return nil, fmt.Errorf("create download file: %s", createErr)
96
		}
97
		tm.MetaInfo = mi
98
		if err := a.cads.Any().GetOrSetMetadata(d.Hex(), &tm); err != nil {
99
			return nil, fmt.Errorf("get or set metainfo: %s", err)
100
		}
101
	} else if err != nil {
102
		return nil, fmt.Errorf("get metainfo: %s", err)
103
	}
104
	t, err := NewTorrent(a.cads, tm.MetaInfo)
105
	if err != nil {
106
		return nil, fmt.Errorf("initialize torrent: %s", err)
107
	}
108
	return t, nil
109
}
110

111
// GetTorrent returns a Torrent for an existing metainfo / file on disk. Ignores namespace.
112
func (a *TorrentArchive) GetTorrent(namespace string, d core.Digest) (storage.Torrent, error) {
113
	var tm metadata.TorrentMeta
114
	if err := a.cads.Any().GetMetadata(d.Hex(), &tm); err != nil {
115
		return nil, fmt.Errorf("get metainfo: %s", err)
116
	}
117
	t, err := NewTorrent(a.cads, tm.MetaInfo)
118
	if err != nil {
119
		return nil, fmt.Errorf("initialize torrent: %s", err)
120
	}
121
	return t, nil
122
}
123

124
// DeleteTorrent deletes a torrent from disk.
125
func (a *TorrentArchive) DeleteTorrent(d core.Digest) error {
126
	if err := a.cads.Any().DeleteFile(d.Hex()); err != nil && !os.IsNotExist(err) {
127
		return err
128
	}
129
	return nil
130
}
131