From 49007969cc72f2bd165d47ca11f6978e1d7f948b Mon Sep 17 00:00:00 2001
From: Abhishek Pandey <abhishek@alcion.ai>
Date: Fri, 30 Jun 2023 02:04:07 -0700
Subject: [PATCH] Add logger extension

---
 src/pkg/extensions/extensions.go      | 55 ++++++++++++++++++++++++---
 src/pkg/extensions/extensions_test.go |  2 +
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/src/pkg/extensions/extensions.go b/src/pkg/extensions/extensions.go
index f130a28f1..9d14c649c 100644
--- a/src/pkg/extensions/extensions.go
+++ b/src/pkg/extensions/extensions.go
@@ -22,6 +22,54 @@ type CorsoItemExtensionFactory func(
 	*details.ExtensionInfo,
 ) (CorsoItemExtension, error)
 
+// Thin wrapper for runtime logging & metrics
+type loggerExtension struct {
+	info    details.ItemInfo
+	innerRc io.ReadCloser
+	ctx     context.Context
+	extInfo *details.ExtensionInfo
+}
+
+func NewMonitoringExtension(
+	ctx context.Context,
+	rc io.ReadCloser,
+	info details.ItemInfo,
+	extInfo *details.ExtensionInfo,
+) (CorsoItemExtension, error) {
+	return &loggerExtension{
+		ctx:     ctx,
+		innerRc: rc,
+		info:    info,
+		extInfo: extInfo,
+	}, nil
+}
+
+func (l *loggerExtension) Read(p []byte) (int, error) {
+	n, err := l.innerRc.Read(p)
+	if err != nil && err != io.EOF {
+		logger.CtxErr(l.ctx, err).Error("inner read")
+		return n, err
+	}
+
+	if err == io.EOF {
+		logger.Ctx(l.ctx).Debug("corso extensions: EOF")
+	}
+
+	return n, err
+}
+
+func (l *loggerExtension) Close() error {
+	err := l.innerRc.Close()
+	if err != nil {
+		logger.CtxErr(l.ctx, err).Error("inner close")
+		return err
+	}
+
+	logger.Ctx(l.ctx).Info("corso extensions: closed")
+
+	return nil
+}
+
 type AddItemExtensioner interface {
 	AddItemExtensions(
 		context.Context,
@@ -43,7 +91,6 @@ func (eh *ItemExtensionHandler) AddItemExtensions(
 	info details.ItemInfo,
 	factories []CorsoItemExtensionFactory,
 ) (io.ReadCloser, *details.ExtensionInfo, error) {
-	// TODO: move to validate
 	if rc == nil {
 		return nil, nil, clues.New("nil inner readcloser")
 	}
@@ -52,14 +99,13 @@ func (eh *ItemExtensionHandler) AddItemExtensions(
 		return nil, nil, clues.New("no extensions supplied")
 	}
 
+	factories = append(factories, NewMonitoringExtension)
 	ctx = clues.Add(ctx, "num_extensions", len(factories))
 
 	extInfo := &details.ExtensionInfo{
 		Data: make(map[string]any),
 	}
 
-	logger.Ctx(ctx).Info("adding extensions")
-
 	for _, factory := range factories {
 		if factory == nil {
 			return nil, nil, clues.New("nil extension factory")
@@ -67,7 +113,7 @@ func (eh *ItemExtensionHandler) AddItemExtensions(
 
 		extRc, err := factory(ctx, rc, info, extInfo)
 		if err != nil {
-			return nil, nil, clues.Wrap(err, "creating extension")
+			return nil, nil, clues.Wrap(err, "calling extension factory")
 		}
 
 		rc = extRc
@@ -75,6 +121,5 @@ func (eh *ItemExtensionHandler) AddItemExtensions(
 
 	logger.Ctx(ctx).Info("added extensions")
 
-	// TODO: Add an outermost extension for logging & metrics
 	return rc, extInfo, nil
 }
diff --git a/src/pkg/extensions/extensions_test.go b/src/pkg/extensions/extensions_test.go
index 6826696d7..49eb28476 100644
--- a/src/pkg/extensions/extensions_test.go
+++ b/src/pkg/extensions/extensions_test.go
@@ -279,6 +279,8 @@ func (suite *ExtensionsUnitSuite) TestAddItemExtensions() {
 	}
 }
 
+// TODO: tests for loggerExtension
+
 // {
 // 	name: "read_failure",
 // 	factories: []CorsoItemExtensionFactory{