From b98b7bdbdd6a1605122adeb66304d38a57983e90 Mon Sep 17 00:00:00 2001 From: Quentin Renard Date: Sun, 5 May 2024 18:58:44 +0200 Subject: [PATCH] Added alloc io context --- class_test.go | 2 +- examples/remuxing/main.go | 2 +- examples/transcoding/main.go | 2 +- format_context_test.go | 2 +- frame_test.go | 7 +- io_context.c | 17 +++ io_context.go | 231 ++++++++++++++++++++++++++++++++++- io_context.h | 9 ++ io_context_test.go | 34 +++++- 9 files changed, 294 insertions(+), 12 deletions(-) create mode 100644 io_context.c create mode 100644 io_context.h diff --git a/class_test.go b/class_test.go index 33ebc41..fc4428c 100644 --- a/class_test.go +++ b/class_test.go @@ -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)) } diff --git a/examples/remuxing/main.go b/examples/remuxing/main.go index 8ef74e5..98765ac 100644 --- a/examples/remuxing/main.go +++ b/examples/remuxing/main.go @@ -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) diff --git a/examples/transcoding/main.go b/examples/transcoding/main.go index 6726856..330fac8 100644 --- a/examples/transcoding/main.go +++ b/examples/transcoding/main.go @@ -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) diff --git a/format_context_test.go b/format_context_test.go index 4485e78..26264c7 100644 --- a/format_context_test.go +++ b/format_context_test.go @@ -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()) diff --git a/frame_test.go b/frame_test.go index dc40796..464e9f4 100644 --- a/frame_test.go +++ b/frame_test.go @@ -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()) diff --git a/io_context.c b/io_context.c new file mode 100644 index 0000000..b948a8f --- /dev/null +++ b/io_context.c @@ -0,0 +1,17 @@ +#include "io_context.h" +#include + +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); +} \ No newline at end of file diff --git a/io_context.go b/io_context.go index eef486a..9049327 100644 --- a/io_context.go +++ b/io_context.go @@ -2,12 +2,20 @@ package astiav //#cgo pkg-config: libavformat //#include +//#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 { @@ -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)) @@ -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) +} diff --git a/io_context.h b/io_context.h new file mode 100644 index 0000000..d77f609 --- /dev/null +++ b/io_context.h @@ -0,0 +1,9 @@ +#include + +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); \ No newline at end of file diff --git a/io_context_test.go b/io_context_test.go index fe4c7f0..2adeea4 100644 --- a/io_context_test.go +++ b/io_context_test.go @@ -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) @@ -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))