Skip to content

Commit

Permalink
Added alloc io context
Browse files Browse the repository at this point in the history
  • Loading branch information
asticode committed May 6, 2024
1 parent 37b487d commit b98b7bd
Show file tree
Hide file tree
Showing 9 changed files with 294 additions and 12 deletions.
2 changes: 1 addition & 1 deletion class_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ func TestClassers(t *testing.T) {
f.Free()
fmc1.Free()
fmc2.CloseInput()
require.NoError(t, ic.Closep())
require.NoError(t, ic.Free())
ssc.Free()
require.Equal(t, cl, len(classers.p))
}
2 changes: 1 addition & 1 deletion examples/remuxing/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ func main() {
if err != nil {
log.Fatal(fmt.Errorf("main: opening io context failed: %w", err))
}
defer ioContext.Closep() //nolint:errcheck
defer ioContext.Free() //nolint:errcheck

// Update output format context
outputFormatContext.SetPb(ioContext)
Expand Down
2 changes: 1 addition & 1 deletion examples/transcoding/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -318,7 +318,7 @@ func openOutputFile() (err error) {
err = fmt.Errorf("main: opening io context failed: %w", err)
return
}
c.AddWithError(ioContext.Closep)
c.AddWithError(ioContext.Free)

// Update output format context
outputFormatContext.SetPb(ioContext)
Expand Down
2 changes: 1 addition & 1 deletion format_context_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ func TestFormatContext(t *testing.T) {
defer fc3.Free()
c, err := OpenIOContext("testdata/video.mp4", NewIOContextFlags(IOContextFlagRead))
require.NoError(t, err)
defer c.Closep() //nolint:errcheck
defer c.Free() //nolint:errcheck
fc3.SetPb(c)
fc3.SetStrictStdCompliance(StrictStdComplianceExperimental)
require.NotNil(t, fc3.Pb())
Expand Down
7 changes: 6 additions & 1 deletion frame_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,12 @@ import (
func TestFrame(t *testing.T) {
f1, err := globalHelper.inputLastFrame("video.mp4", MediaTypeVideo)
require.NoError(t, err)
require.Equal(t, [8]int{384, 192, 192, 0, 0, 0, 0, 0}, f1.Linesize())
// Should be "{384, 192, 192, 0, 0, 0, 0, 0}" but for some reason it"s "{320, 160, 160, 0, 0, 0, 0, 0}"
// on darwin when testing using github
require.Contains(t, [][8]int{
{384, 192, 192, 0, 0, 0, 0, 0},
{320, 160, 160, 0, 0, 0, 0, 0},
}, f1.Linesize())
require.Equal(t, int64(60928), f1.PktDts())
require.Equal(t, unsafe.Pointer(f1.c), f1.UnsafePointer())

Expand Down
17 changes: 17 additions & 0 deletions io_context.c
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
#include "io_context.h"
#include <stdint.h>

int astiavIOContextReadFunc(void *opaque, uint8_t *buf, int buf_size)
{
return goAstiavIOContextReadFunc(opaque, buf, buf_size);
}

int64_t astiavIOContextSeekFunc(void *opaque, int64_t offset, int whence)
{
return goAstiavIOContextSeekFunc(opaque, offset, whence);
}

int astiavIOContextWriteFunc(void *opaque, uint8_t *buf, int buf_size)
{
return goAstiavIOContextWriteFunc(opaque, buf, buf_size);
}
231 changes: 226 additions & 5 deletions io_context.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,20 @@ package astiav

//#cgo pkg-config: libavformat
//#include <libavformat/avformat.h>
//#include "io_context.h"
import "C"
import "unsafe"
import (
"errors"
"fmt"
"sync"
"unsafe"
)

// https://github.com/FFmpeg/FFmpeg/blob/n5.0/libavformat/avio.h#L161
type IOContext struct {
c *C.struct_AVIOContext
buffer unsafe.Pointer
c *C.struct_AVIOContext
handlerID unsafe.Pointer
}

func newIOContextFromC(c *C.struct_AVIOContext) *IOContext {
Expand All @@ -21,6 +29,80 @@ func newIOContextFromC(c *C.struct_AVIOContext) *IOContext {

var _ Classer = (*IOContext)(nil)

type IOContextReadFunc func(b []byte) (n int, err error)

type IOContextSeekFunc func(offset int64, whence int) (n int64, err error)

type IOContextWriteFunc func(b []byte) (n int, err error)

func AllocIOContext(bufferSize int, readFunc IOContextReadFunc, seekFunc IOContextSeekFunc, writeFunc IOContextWriteFunc) (ic *IOContext, err error) {
// Invalid buffer size
if bufferSize <= 0 {
err = errors.New("astiav: buffer size <= 0")
return
}

// Alloc buffer
buffer := C.av_malloc(C.size_t(bufferSize))
if buffer == nil {
err = errors.New("astiav: allocating buffer failed")
return
}

// Make sure buffer is freed in case of error
defer func() {
if err != nil {
C.av_free(buffer)
}
}()

// Since go doesn't allow c to store pointers to go data, we need to create this C pointer
handlerID := C.malloc(C.size_t(1))
if handlerID == nil {
err = errors.New("astiav: allocating handler id failed")
return
}

// Make sure handler id is freed in case of error
defer func() {
if err != nil {
C.free(handlerID)
}
}()

// Get callbacks
var cReadFunc, cSeekFunc, cWriteFunc *[0]byte
if readFunc != nil {
cReadFunc = (*[0]byte)(C.astiavIOContextReadFunc)
}
if seekFunc != nil {
cSeekFunc = (*[0]byte)(C.astiavIOContextSeekFunc)
}
if writeFunc != nil {
cWriteFunc = (*[0]byte)(C.astiavIOContextWriteFunc)
}

// Alloc io context
cic := C.avio_alloc_context((*C.uchar)(buffer), C.int(bufferSize), 1, handlerID, cReadFunc, cWriteFunc, cSeekFunc)
if cic == nil {
err = errors.New("astiav: allocating io context failed: %w")
return
}

// Create io context
ic = newIOContextFromC(cic)

// Store buffer and handler
ic.buffer = buffer
ic.handlerID = handlerID
ioContextHandlers.set(handlerID, &ioContextHandler{
r: readFunc,
s: seekFunc,
w: writeFunc,
})
return
}

func OpenIOContext(filename string, flags IOContextFlags) (*IOContext, error) {
cfi := C.CString(filename)
defer C.free(unsafe.Pointer(cfi))
Expand All @@ -35,17 +117,156 @@ func (ic *IOContext) Class() *Class {
return newClassFromC(unsafe.Pointer(ic.c))
}

func (ic *IOContext) Closep() error {
func (ic *IOContext) Free() error {
classers.del(ic)
if ic.c != nil {
return newError(C.avio_closep(&ic.c))
if ic.buffer == nil {
if err := newError(C.avio_closep(&ic.c)); err != nil {
return err
}
} else {
C.av_free(ic.buffer)
ic.buffer = nil
C.free(ic.handlerID)
ic.handlerID = nil
C.avio_context_free(&ic.c)
ic.c = nil
}
}
return nil
}

func (ic *IOContext) Read(b []byte) (n int, err error) {
// Nothing to read
if b == nil || len(b) <= 0 {
return
}

// Alloc buffer
buf := C.av_malloc(C.size_t(len(b)))
if buf == nil {
err = errors.New("astiav: allocating buffer failed")
return
}

// Make sure buffer is freed
defer C.av_free(buf)

// Read
ret := C.avio_read_partial(ic.c, (*C.uchar)(unsafe.Pointer(buf)), C.int(len(b)))
if err = newError(ret); err != nil {
err = fmt.Errorf("astiav: reading failed: %w", err)
return
}

// Copy
C.memcpy(unsafe.Pointer(&b[0]), unsafe.Pointer(buf), C.size_t(ret))
n = int(ret)
return
}

func (ic *IOContext) Seek(offset int64, whence int) (int64, error) {
ret := C.avio_seek(ic.c, C.int64_t(offset), C.int(whence))
if err := newError(C.int(ret)); err != nil {
return 0, err
}
return int64(ret), nil
}

func (ic *IOContext) Write(b []byte) {
if b == nil {
// Nothing to write
if b == nil || len(b) <= 0 {
return
}

// Write
C.avio_write(ic.c, (*C.uchar)(unsafe.Pointer(&b[0])), C.int(len(b)))
}

func (ic *IOContext) Flush() {
C.avio_flush(ic.c)
}

type ioContextHandler struct {
r IOContextReadFunc
s IOContextSeekFunc
w IOContextWriteFunc
}

var ioContextHandlers = newIOContextHandlerPool()

type ioContextHandlerPool struct {
m sync.Mutex
p map[unsafe.Pointer]*ioContextHandler
}

func newIOContextHandlerPool() *ioContextHandlerPool {
return &ioContextHandlerPool{p: make(map[unsafe.Pointer]*ioContextHandler)}
}

func (p *ioContextHandlerPool) set(id unsafe.Pointer, h *ioContextHandler) {
p.m.Lock()
defer p.m.Unlock()
p.p[id] = h
}

func (p *ioContextHandlerPool) get(id unsafe.Pointer) (h *ioContextHandler, ok bool) {
p.m.Lock()
defer p.m.Unlock()
h, ok = p.p[id]
return
}

//export goAstiavIOContextReadFunc
func goAstiavIOContextReadFunc(opaque unsafe.Pointer, buf *C.uint8_t, bufSize C.int) C.int {
// Get handler
h, ok := ioContextHandlers.get(opaque)
if !ok {
return C.AVERROR_UNKNOWN
}

// Create go buffer
b := make([]byte, int(bufSize), int(bufSize))

// Read
n, err := h.r(b)
if err != nil {
return C.AVERROR_UNKNOWN
}

// Copy
C.memcpy(unsafe.Pointer(buf), unsafe.Pointer(&b[0]), C.size_t(n))
return C.int(n)
}

//export goAstiavIOContextSeekFunc
func goAstiavIOContextSeekFunc(opaque unsafe.Pointer, offset C.int64_t, whence C.int) C.int64_t {
// Get handler
h, ok := ioContextHandlers.get(opaque)
if !ok {
return C.AVERROR_UNKNOWN
}

// Seek
n, err := h.s(int64(offset), int(whence))
if err != nil {
return C.int64_t(C.AVERROR_UNKNOWN)
}
return C.int64_t(n)
}

//export goAstiavIOContextWriteFunc
func goAstiavIOContextWriteFunc(opaque unsafe.Pointer, buf *C.uint8_t, bufSize C.int) C.int {
// Get handler
h, ok := ioContextHandlers.get(opaque)
if !ok {
return C.AVERROR_UNKNOWN
}

// Write
n, err := h.w(C.GoBytes(unsafe.Pointer(buf), bufSize))
if err != nil {
return C.AVERROR_UNKNOWN
}
return C.int(n)
}
9 changes: 9 additions & 0 deletions io_context.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
#include <stdint.h>

extern int goAstiavIOContextReadFunc(void *opaque, uint8_t *buf, int buf_size);
extern int64_t goAstiavIOContextSeekFunc(void *opaque, int64_t offset, int whence);
extern int goAstiavIOContextWriteFunc(void *opaque, uint8_t *buf, int buf_size);

int astiavIOContextReadFunc(void *opaque, uint8_t *buf, int buf_size);
int64_t astiavIOContextSeekFunc(void *opaque, int64_t offset, int whence);
int astiavIOContextWriteFunc(void *opaque, uint8_t *buf, int buf_size);
34 changes: 32 additions & 2 deletions io_context_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,37 @@ import (
)

func TestIOContext(t *testing.T) {
var seeked bool
rb := []byte("read")
wb := []byte("write")
var written []byte
c, err := AllocIOContext(8, func(b []byte) (int, error) {
copy(b, rb)
return len(rb), nil
}, func(offset int64, whence int) (n int64, err error) {
seeked = true
return offset, nil
}, func(b []byte) (int, error) {
written = make([]byte, len(b))
copy(written, b)
return len(b), nil
})
require.NoError(t, err)
defer c.Free()
b := make([]byte, 6)
n, err := c.Read(b)
require.NoError(t, err)
require.Equal(t, 4, n)
require.Equal(t, rb, b[:n])
_, err = c.Seek(2, 0)
require.NoError(t, err)
require.True(t, seeked)
c.Write(wb)
c.Flush()
require.Equal(t, wb, written)
}

func TestOpenIOContext(t *testing.T) {
path := filepath.Join(t.TempDir(), "iocontext.txt")
c, err := OpenIOContext(path, NewIOContextFlags(IOContextFlagWrite))
require.NoError(t, err)
Expand All @@ -17,8 +48,7 @@ func TestIOContext(t *testing.T) {
require.Equal(t, "AVIOContext", cl.Name())
c.Write(nil)
c.Write([]byte("test"))
err = c.Closep()
require.NoError(t, err)
require.NoError(t, c.Free())
b, err := os.ReadFile(path)
require.NoError(t, err)
require.Equal(t, "test", string(b))
Expand Down

0 comments on commit b98b7bd

Please sign in to comment.