-
Notifications
You must be signed in to change notification settings - Fork 16
/
disk_mate_shard.go
218 lines (194 loc) · 5.98 KB
/
disk_mate_shard.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
package bampair
import (
"bytes"
"encoding/binary"
"fmt"
"io"
"os"
"path"
"sync"
"time"
"github.com/golang/snappy"
"github.com/grailbio/base/log"
grailbam "github.com/grailbio/bio/encoding/bam"
"github.com/grailbio/hts/bam"
"github.com/grailbio/hts/sam"
)
// diskMateShard is an on-disk shard of the distant mates. It is
// meant to store a portion of the distant mates in a way that allows
// the caller to load the distant mates into memory only when they are
// useful.
//
// The number of diskMateShards is controlled by a command line
// parameter passed through distantMateTable. A reasonable number is
// 1000 diskMateShards, which results in the largest shards being
// roughly 200MB for a 130GB input file with 5% distant mates. The
// diskMateShards are not evenly sized because the input file is
// unevenly sharded, and has uneven distribution of distant mates.
// The input bam file's shards map to these disk shards evenly, so an
// input bam has 30,000 shards, then the first 30 input shards map to
// the first diskMateShard, and the next 30 input shards map to the
// second diskMateShard.
//
// add(), openReader() and closeReader() are threadsafe. The call to
// closeWriter() should occur after all calls to add() are finished.
// Calls to getMate() should occur only between a call to openReader()
// and a call to closeReader().
type diskMateShard struct {
header *sam.Header
f *os.File
writer io.WriteCloser
m map[string][]*indexedRecord
mateShardIdx int
mutex sync.Mutex
refcount int
buf bytes.Buffer
}
// newDiskMateShard returns a new diskMateShard.
func newDiskMateShard(header *sam.Header, tempDir string, mateShardIdx, numMateShards int) (*diskMateShard, error) {
filename := path.Join(tempDir, fmt.Sprintf("mates_%04d_of_%04d", mateShardIdx, numMateShards))
f, err := os.Create(filename)
if err != nil {
return nil, fmt.Errorf("error creating file %s: %v", filename, err)
}
return &diskMateShard{
header: header,
f: f,
writer: snappy.NewBufferedWriter(f),
mateShardIdx: mateShardIdx,
}, nil
}
func (s *diskMateShard) add(mate *sam.Record, fileIdx uint64) error {
s.mutex.Lock()
defer s.mutex.Unlock()
if err := binary.Write(s.writer, binary.LittleEndian, fileIdx); err != nil {
return fmt.Errorf("error writing fileIdx to mate shard: %v", err)
}
s.buf.Reset()
if err := bam.Marshal(mate, &s.buf); err != nil {
return fmt.Errorf("error marshalling to mate shard: %v", err)
}
if _, err := s.writer.Write(s.buf.Bytes()); err != nil {
return fmt.Errorf("error writing record to mate shard: %v", err)
}
return nil
}
func (s *diskMateShard) closeWriter() error {
s.mutex.Lock()
defer s.mutex.Unlock()
if err := s.writer.Close(); err != nil {
return fmt.Errorf("failed to close snappy writer for distantMateShard %s: %v", s.f.Name(), err)
}
if err := s.f.Close(); err != nil {
return fmt.Errorf("failed to close f for distantMateShard %s: %v", s.f.Name(), err)
}
return nil
}
func (s *diskMateShard) openReader() error {
s.mutex.Lock()
defer s.mutex.Unlock()
s.refcount++
if s.refcount > 1 {
return nil
}
f, err := os.Open(s.f.Name())
if err != nil {
return fmt.Errorf("error opening file %s: %v", s.f.Name(), err)
}
reader := snappy.NewReader(f)
t0 := time.Now()
s.m = map[string][]*indexedRecord{}
records := 0
log.Debug.Printf("reading disk shard %d", s.mateShardIdx)
buf := make([]byte, 1024)
for {
// Read fileIdx and record size.
n, err := io.ReadFull(reader, buf[0:12])
if n > 0 && n < 12 {
return fmt.Errorf("Not enough bytes: %d, file %s", n, s.f.Name())
}
if err == io.EOF {
break
} else if err != nil {
return fmt.Errorf("Error reading mate shard: %v", err)
}
fileIdx := binary.LittleEndian.Uint64(buf[0:8])
size := binary.LittleEndian.Uint32(buf[8:12])
// Grow the slice if necessary
for uint32(cap(buf)) < size {
buf = append(buf[:cap(buf)], 0)
}
buf = buf[:size]
n, err = io.ReadFull(reader, buf)
if uint32(n) < size {
return fmt.Errorf("Not enough bytes: %d expected %d file %s", n, size, s.f.Name())
}
if err != nil && err != io.EOF {
return fmt.Errorf("Error reading mate shard: %v", err)
}
// Unmarshal record
record, err := grailbam.Unmarshal(buf, s.header)
if err != nil {
return fmt.Errorf("Failed to unmarshal disk mate: %v", err)
}
records++
iRecord := &indexedRecord{record, fileIdx}
e, found := s.m[iRecord.r.Name]
if found {
// We might receive duplicate records here because
// multiple inputs shards map to one mateShardIdx. Add
// only unique records.
duplicate := false
for _, existing := range e {
if recordsEqual(existing.r, iRecord.r) {
duplicate = true
break
}
}
if duplicate {
continue
}
// There's no particular order for left and right in a distantMateEntry.
s.m[iRecord.r.Name] = append(s.m[iRecord.r.Name], iRecord)
if len(s.m[iRecord.r.Name]) > 2 {
return fmt.Errorf("Got 3rd read, %v %d", iRecord.r, iRecord.fileIdx)
}
} else {
s.m[iRecord.r.Name] = []*indexedRecord{iRecord}
}
}
if err := f.Close(); err != nil {
return fmt.Errorf("Error closing %s: %v", s.f.Name(), err)
}
t1 := time.Now()
log.Debug.Printf("read disk shard %d with %d records in %v", s.mateShardIdx, records, t1.Sub(t0))
return nil
}
func (s *diskMateShard) getMate(shardInfo *ShardInfo, r *sam.Record) (*sam.Record, uint64) {
e := s.m[r.Name]
for _, iRecord := range e {
if !recordsEqual(iRecord.r, r) {
return iRecord.r, iRecord.fileIdx + shardInfo.getShardStartFileIdx(iRecord.r)
}
}
return nil, 0
}
func (s *diskMateShard) closeReader() {
s.mutex.Lock()
defer s.mutex.Unlock()
s.refcount--
// Zero out the map to free the memory.
if s.refcount == 0 {
s.m = nil
}
}
func recordsEqual(a, b *sam.Record) bool {
return a.Name == b.Name &&
a.Ref.ID() == b.Ref.ID() &&
a.Pos == b.Pos &&
a.MapQ == b.MapQ &&
a.Flags == b.Flags &&
a.MateRef.ID() == b.MateRef.ID() &&
a.MatePos == b.MatePos &&
a.TempLen == b.TempLen
}