storage/mmap.go

   1 package storage
   2
   3 import (
   4         "fmt"
   5         "io"
   6         "os"
   7         "path/filepath"
   8
   9         "github.com/anacrolix/missinggo"
  10         "github.com/edsrzf/mmap-go"
  11
  12         "github.com/anacrolix/torrent/metainfo"
  13         "github.com/anacrolix/torrent/mmap_span"
  14 )
  15
  16 type mmapStorage struct {
  17         baseDir string
  18         pc      PieceCompletion
  19 }
  20
  21 func NewMMap(baseDir string) ClientImpl {
  22         return NewMMapWithCompletion(baseDir, pieceCompletionForDir(baseDir))
  23 }
  24
  25 func NewMMapWithCompletion(baseDir string, completion PieceCompletion) ClientImpl {
  26         return &mmapStorage{
  27                 baseDir: baseDir,
  28                 pc:      completion,
  29         }
  30 }
  31
  32 func (s *mmapStorage) OpenTorrent(info *metainfo.Info, infoHash metainfo.Hash) (t TorrentImpl, err error) {
  33         span, err := mMapTorrent(info, s.baseDir)
  34         t = &mmapTorrentStorage{
  35                 span: span,
  36                 pc:   s.pc,
  37         }
  38         return
  39 }
  40
  41 func (s *mmapStorage) Close() error {
  42         return s.pc.Close()
  43 }
  44
  45 type mmapTorrentStorage struct {
  46         span mmap_span.MMapSpan
  47         pc   PieceCompletion
  48 }
  49
  50 func (ts *mmapTorrentStorage) Piece(p metainfo.Piece) PieceImpl {
  51         return mmapStoragePiece{
  52                 pc:       ts.pc,
  53                 p:        p,
  54                 ReaderAt: io.NewSectionReader(ts.span, p.Offset(), p.Length()),
  55                 WriterAt: missinggo.NewSectionWriter(ts.span, p.Offset(), p.Length()),
  56         }
  57 }
  58
  59 func (ts *mmapTorrentStorage) Close() error {
  60         ts.pc.Close()
  61         return ts.span.Close()
  62 }
  63
  64 type mmapStoragePiece struct {
  65         pc PieceCompletion
  66         p  metainfo.Piece
  67         ih metainfo.Hash
  68         io.ReaderAt
  69         io.WriterAt
  70 }
  71
  72 func (me mmapStoragePiece) pieceKey() metainfo.PieceKey {
  73         return metainfo.PieceKey{me.ih, me.p.Index()}
  74 }
  75
  76 func (sp mmapStoragePiece) GetIsComplete() (ret bool) {
  77         ret, _ = sp.pc.Get(sp.pieceKey())
  78         return
  79 }
  80
  81 func (sp mmapStoragePiece) MarkComplete() error {
  82         sp.pc.Set(sp.pieceKey(), true)
  83         return nil
  84 }
  85
  86 func (sp mmapStoragePiece) MarkNotComplete() error {
  87         sp.pc.Set(sp.pieceKey(), false)
  88         return nil
  89 }
  90
  91 func mMapTorrent(md *metainfo.Info, location string) (mms mmap_span.MMapSpan, err error) {
  92         defer func() {
  93                 if err != nil {
  94                         mms.Close()
  95                 }
  96         }()
  97         for _, miFile := range md.UpvertedFiles() {
  98                 fileName := filepath.Join(append([]string{location, md.Name}, miFile.Path...)...)
  99                 err = os.MkdirAll(filepath.Dir(fileName), 0777)
 100                 if err != nil {
 101                         err = fmt.Errorf("error creating data directory %q: %s", filepath.Dir(fileName), err)
 102                         return
 103                 }
 104                 var file *os.File
 105                 file, err = os.OpenFile(fileName, os.O_CREATE|os.O_RDWR, 0666)
 106                 if err != nil {
 107                         return
 108                 }
 109                 func() {
 110                         defer file.Close()
 111                         var fi os.FileInfo
 112                         fi, err = file.Stat()
 113                         if err != nil {
 114                                 return
 115                         }
 116                         if fi.Size() < miFile.Length {
 117                                 err = file.Truncate(miFile.Length)
 118                                 if err != nil {
 119                                         return
 120                                 }
 121                         }
 122                         if miFile.Length == 0 {
 123                                 // Can't mmap() regions with length 0.
 124                                 return
 125                         }
 126                         var mMap mmap.MMap
 127                         mMap, err = mmap.MapRegion(file,
 128                                 int(miFile.Length), // Probably not great on <64 bit systems.
 129                                 mmap.RDWR, 0, 0)
 130                         if err != nil {
 131                                 err = fmt.Errorf("error mapping file %q, length %d: %s", file.Name(), miFile.Length, err)
 132                                 return
 133                         }
 134                         if int64(len(mMap)) != miFile.Length {
 135                                 panic("mmap has wrong length")
 136                         }
 137                         mms.Append(mMap)
 138                 }()
 139                 if err != nil {
 140                         return
 141                 }
 142         }
 143         return
 144 }