Add transient storage lifecycle receiver framework

Author: xiaokangwang

common/environment/rootcap_impl.go

@@ -74,7 +74,8 @@ func (r *rootEnvImpl) DropProxyEnvironment(tag string) error {
 	if err != nil {
 		return err
 	}
-	return transientStorage.DropScope(r.ctx, tag)
+	transientStorage.Clear(r.ctx)
+	return r.transientStorage.DropScope(r.ctx, tag)
 }
 
 type appEnvImpl struct {

common/environment/transientstorageimpl/storage.go

@@ -56,7 +56,16 @@ func (s *scopedTransientStorageImpl) List(ctx context.Context, keyPrefix string)
 func (s *scopedTransientStorageImpl) Clear(ctx context.Context) {
 	s.access.Lock()
 	defer s.access.Unlock()
+	for _, v := range s.values {
+		if sw, ok := v.(storage.TransientStorageLifecycleReceiver); ok {
+			_ = sw.Close()
+		}
+	}
 	s.values = map[string]interface{}{}
+	for _, v := range s.scopes {
+		v.Clear(ctx)
+	}
+	s.scopes = map[string]storage.ScopedTransientStorage{}
 }
 
 func (s *scopedTransientStorageImpl) NarrowScope(ctx context.Context, key string) (storage.ScopedTransientStorage, error) {
@@ -74,6 +83,9 @@ func (s *scopedTransientStorageImpl) NarrowScope(ctx context.Context, key string
 func (s *scopedTransientStorageImpl) DropScope(ctx context.Context, key string) error {
 	s.access.Lock()
 	defer s.access.Unlock()
+	if v, ok := s.scopes[key]; ok {
+		v.Clear(ctx)
+	}
 	delete(s.scopes, key)
 	return nil
 }

features/extension/storage/storage.go

@@ -3,6 +3,7 @@ package storage
 import (
 	"context"
 
+	"github.com/v2fly/v2ray-core/v5/common"
 	"github.com/v2fly/v2ray-core/v5/features"
 )
 
@@ -32,3 +33,8 @@ type ScopedPersistentStorageService interface {
 }
 
 var ScopedPersistentStorageServiceType = (*ScopedPersistentStorageService)(nil)
+
+type TransientStorageLifecycleReceiver interface {
+	IsTransientStorageLifecycleReceiver()
+	common.Closable
+}

transport/internet/grpc/dial.go

@@ -17,6 +17,8 @@ import (
 
 	core "github.com/v2fly/v2ray-core/v5"
 	"github.com/v2fly/v2ray-core/v5/common"
+	"github.com/v2fly/v2ray-core/v5/common/environment"
+	"github.com/v2fly/v2ray-core/v5/common/environment/envctx"
 	"github.com/v2fly/v2ray-core/v5/common/net"
 	"github.com/v2fly/v2ray-core/v5/common/session"
 	"github.com/v2fly/v2ray-core/v5/transport/internet"
@@ -38,12 +40,25 @@ func init() {
 	common.Must(internet.RegisterTransportDialer(protocolName, Dial))
 }
 
-type dialerCanceller func()
+type transportConnectionState struct {
+	scopedDialerMap    map[net.Destination]*grpc.ClientConn
+	scopedDialerAccess sync.Mutex
+}
 
-var (
-	globalDialerMap    map[net.Destination]*grpc.ClientConn
-	globalDialerAccess sync.Mutex
-)
+func (t *transportConnectionState) IsTransientStorageLifecycleReceiver() {
+}
+
+func (t *transportConnectionState) Close() error {
+	t.scopedDialerAccess.Lock()
+	defer t.scopedDialerAccess.Unlock()
+	for _, conn := range t.scopedDialerMap {
+		_ = conn.Close()
+	}
+	t.scopedDialerMap = nil
+	return nil
+}
+
+type dialerCanceller func()
 
 func dialgRPC(ctx context.Context, dest net.Destination, streamSettings *internet.MemoryStreamConfig) (net.Conn, error) {
 	grpcSettings := streamSettings.ProtocolSettings.(*Config)
@@ -70,25 +85,36 @@ func dialgRPC(ctx context.Context, dest net.Destination, streamSettings *interne
 }
 
 func getGrpcClient(ctx context.Context, dest net.Destination, dialOption grpc.DialOption, streamSettings *internet.MemoryStreamConfig) (*grpc.ClientConn, dialerCanceller, error) {
-	globalDialerAccess.Lock()
-	defer globalDialerAccess.Unlock()
+	transportEnvironment := envctx.EnvironmentFromContext(ctx).(environment.TransportEnvironment)
+	state, err := transportEnvironment.TransientStorage().Get(ctx, "grpc-transport-connection-state")
+	if err != nil {
+		state = &transportConnectionState{}
+		transportEnvironment.TransientStorage().Put(ctx, "grpc-transport-connection-state", state)
+		state, err = transportEnvironment.TransientStorage().Get(ctx, "grpc-transport-connection-state")
+		if err != nil {
+			return nil, nil, newError("failed to get grpc transport connection state").Base(err)
+		}
+	}
+	stateTyped := state.(*transportConnectionState)
+
+	stateTyped.scopedDialerAccess.Lock()
+	defer stateTyped.scopedDialerAccess.Unlock()
 
-	if globalDialerMap == nil {
-		globalDialerMap = make(map[net.Destination]*grpc.ClientConn)
+	if stateTyped.scopedDialerMap == nil {
+		stateTyped.scopedDialerMap = make(map[net.Destination]*grpc.ClientConn)
 	}
 
 	canceller := func() {
-		globalDialerAccess.Lock()
-		defer globalDialerAccess.Unlock()
-		delete(globalDialerMap, dest)
+		stateTyped.scopedDialerAccess.Lock()
+		defer stateTyped.scopedDialerAccess.Unlock()
+		delete(stateTyped.scopedDialerMap, dest)
 	}
 
-	// TODO Should support chain proxy to the same destination
-	if client, found := globalDialerMap[dest]; found && client.GetState() != connectivity.Shutdown {
+	if client, found := stateTyped.scopedDialerMap[dest]; found && client.GetState() != connectivity.Shutdown {
 		return client, canceller, nil
 	}
 
-	conn, err := grpc.Dial(
+	conn, err := grpc.NewClient(
 		dest.Address.String()+":"+dest.Port.String(),
 		dialOption,
 		grpc.WithConnectParams(grpc.ConnectParams{
@@ -117,6 +143,14 @@ func getGrpcClient(ctx context.Context, dest net.Destination, dialOption grpc.Di
 			return internet.DialSystem(detachedContext, net.TCPDestination(address, port), streamSettings.SocketSettings)
 		}),
 	)
-	globalDialerMap[dest] = conn
+	canceller = func() {
+		stateTyped.scopedDialerAccess.Lock()
+		defer stateTyped.scopedDialerAccess.Unlock()
+		delete(stateTyped.scopedDialerMap, dest)
+		if err != nil {
+			conn.Close()
+		}
+	}
+	stateTyped.scopedDialerMap[dest] = conn
 	return conn, canceller, err
 }