ooni-probe-cli/internal/engine/experiment/ndt7/upload.go
Simone Basso d57c78bc71
chore: merge probe-engine into probe-cli (#201)
This is how I did it:

1. `git clone https://github.com/ooni/probe-engine internal/engine`

2. ```
(cd internal/engine && git describe --tags)
v0.23.0
```

3. `nvim go.mod` (merging `go.mod` with `internal/engine/go.mod`

4. `rm -rf internal/.git internal/engine/go.{mod,sum}`

5. `git add internal/engine`

6. `find . -type f -name \*.go -exec sed -i 's@/ooni/probe-engine@/ooni/probe-cli/v3/internal/engine@g' {} \;`

7. `go build ./...` (passes)

8. `go test -race ./...` (temporary failure on RiseupVPN)

9. `go mod tidy`

10. this commit message

Once this piece of work is done, we can build a new version of `ooniprobe` that
is using `internal/engine` directly. We need to do more work to ensure all the
other functionality in `probe-engine` (e.g. making mobile packages) are still WAI.

Part of https://github.com/ooni/probe/issues/1335
2021-02-02 12:05:47 +01:00

76 lines
1.8 KiB
Go

package ndt7
import (
"context"
"time"
"github.com/gorilla/websocket"
)
func newMessage(n int) (*websocket.PreparedMessage, error) {
return websocket.NewPreparedMessage(websocket.BinaryMessage, make([]byte, n))
}
type uploadManager struct {
conn mockableConn
fractionForScaling int64
maxRuntime time.Duration
maxMessageSize int
maxScaledMessageSize int
measureInterval time.Duration
minMessageSize int
newMessage func(int) (*websocket.PreparedMessage, error)
onPerformance callbackPerformance
}
func newUploadManager(
conn mockableConn, onPerformance callbackPerformance,
) uploadManager {
return uploadManager{
conn: conn,
fractionForScaling: paramFractionForScaling,
maxRuntime: paramMaxRuntime,
maxMessageSize: paramMaxMessageSize,
maxScaledMessageSize: paramMaxScaledMessageSize,
measureInterval: paramMeasureInterval,
minMessageSize: paramMinMessageSize,
newMessage: newMessage,
onPerformance: onPerformance,
}
}
func (mgr uploadManager) run(ctx context.Context) error {
var total int64
start := time.Now()
if err := mgr.conn.SetWriteDeadline(time.Now().Add(mgr.maxRuntime)); err != nil {
return err
}
size := mgr.minMessageSize
message, err := mgr.newMessage(size)
if err != nil {
return err
}
ticker := time.NewTicker(mgr.measureInterval)
defer ticker.Stop()
for ctx.Err() == nil {
if err := mgr.conn.WritePreparedMessage(message); err != nil {
return err
}
total += int64(size)
select {
case now := <-ticker.C:
mgr.onPerformance(now.Sub(start), total)
default:
// NOTHING
}
if size >= mgr.maxScaledMessageSize || int64(size) >= (total/mgr.fractionForScaling) {
continue
}
size <<= 1
if message, err = mgr.newMessage(size); err != nil {
return err
}
}
return nil
}