Fix sourceJar task dependencies (#28900)

diff --git a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
index 9f4ae8b..705bcb9 100644
--- a/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
+++ b/buildSrc/src/main/groovy/org/apache/beam/gradle/BeamModulePlugin.groovy
@@ -2395,7 +2395,20 @@
 
     // TODO: Decide whether this should be inlined into the one project that relies on it
     // or be left here.
-    project.ext.applyAvroNature = { project.apply plugin: "com.commercehub.gradle.plugin.avro" }
+    project.ext.applyAvroNature = {
+      project.apply plugin: "com.commercehub.gradle.plugin.avro"
+
+      // add dependency BeamModulePlugin defined custom tasks
+      // they are defined only when certain flags are provided (e.g. -Prelease; -Ppublishing, etc)
+      def sourcesJar = project.tasks.findByName('sourcesJar')
+      if (sourcesJar != null) {
+        sourcesJar.dependsOn project.tasks.getByName('generateAvroJava')
+      }
+      def testSourcesJar = project.tasks.findByName('testSourcesJar')
+      if (testSourcesJar != null) {
+        testSourcesJar.dependsOn project.tasks.getByName('generateTestAvroJava')
+      }
+    }
 
     project.ext.applyAntlrNature = {
       project.apply plugin: 'antlr'
@@ -2406,6 +2419,17 @@
           generatedSourceDirs += project.generateTestGrammarSource.outputDirectory
         }
       }
+
+      // add dependency BeamModulePlugin defined custom tasks
+      // they are defined only when certain flags are provided (e.g. -Prelease; -Ppublishing, etc)
+      def sourcesJar = project.tasks.findByName('sourcesJar')
+      if (sourcesJar != null) {
+        sourcesJar.mustRunAfter project.tasks.getByName('generateGrammarSource')
+      }
+      def testSourcesJar = project.tasks.findByName('testSourcesJar')
+      if (testSourcesJar != null) {
+        testSourcesJar.dependsOn project.tasks.getByName('generateTestGrammarSource')
+      }
     }
 
     // Creates a task to run the quickstart for a runner.
diff --git a/runners/flink/flink_runner.gradle b/runners/flink/flink_runner.gradle
index 30fb922..c087575 100644
--- a/runners/flink/flink_runner.gradle
+++ b/runners/flink/flink_runner.gradle
@@ -46,36 +46,51 @@
 /*
  * Copy & merge source overrides into build directory.
  */
-def sourceOverridesBase = "${project.buildDir}/source-overrides/src"
+def sourceOverridesBase = project.layout.buildDirectory.dir('source-overrides/src')
 
 def copySourceOverrides = tasks.register('copySourceOverrides', Copy) {
   it.from main_source_overrides
   it.into "${sourceOverridesBase}/main/java"
   it.duplicatesStrategy DuplicatesStrategy.INCLUDE
 }
-compileJava.dependsOn copySourceOverrides
 
 def copyResourcesOverrides = tasks.register('copyResourcesOverrides', Copy) {
   it.from main_resources_overrides
   it.into "${sourceOverridesBase}/main/resources"
   it.duplicatesStrategy DuplicatesStrategy.INCLUDE
 }
-processResources.dependsOn copyResourcesOverrides
 
 def copyTestSourceOverrides = tasks.register('copyTestSourceOverrides', Copy) {
   it.from test_source_overrides
   it.into "${sourceOverridesBase}/test/java"
   it.duplicatesStrategy DuplicatesStrategy.INCLUDE
 }
-compileTestJava.dependsOn copyTestSourceOverrides
 
 def copyTestResourcesOverrides = tasks.register('copyTestResourcesOverrides', Copy) {
   it.from test_resources_overrides
   it.into "${sourceOverridesBase}/test/resources"
   it.duplicatesStrategy DuplicatesStrategy.INCLUDE
 }
+
+// add dependency to gradle Java plugin defined tasks
+compileJava.dependsOn copySourceOverrides
+processResources.dependsOn copyResourcesOverrides
+compileTestJava.dependsOn copyTestSourceOverrides
 processTestResources.dependsOn copyTestResourcesOverrides
 
+// add dependency BeamModulePlugin defined custom tasks
+// they are defined only when certain flags are provided (e.g. -Prelease; -Ppublishing, etc)
+def sourcesJar = project.tasks.findByName('sourcesJar')
+if (sourcesJar != null) {
+  sourcesJar.dependsOn copySourceOverrides
+  sourcesJar.dependsOn copyResourcesOverrides
+}
+def testSourcesJar = project.tasks.findByName('testSourcesJar')
+if (testSourcesJar != null) {
+  testSourcesJar.dependsOn copyTestSourceOverrides
+  testSourcesJar.dependsOn copyTestResourcesOverrides
+}
+
 /*
  * We have to explicitly set all directories here to make sure each
  * version of Flink has the correct overrides set.
diff --git a/sdks/java/maven-archetypes/examples/build.gradle b/sdks/java/maven-archetypes/examples/build.gradle
index 56b4a7c..1edb55a 100644
--- a/sdks/java/maven-archetypes/examples/build.gradle
+++ b/sdks/java/maven-archetypes/examples/build.gradle
@@ -72,6 +72,13 @@
   commandLine './generate-sources.sh'
 }
 
+// add dependency BeamModulePlugin defined custom tasks
+// they are defined only when certain flags are provided (e.g. -Prelease; -Ppublishing, etc)
+def sourcesJar = project.tasks.findByName('sourcesJar')
+if (sourcesJar != null) {
+  sourcesJar.dependsOn generateSources
+}
+
 sourceSets {
  main {
   output.dir('src', builtBy: 'generateSources')
diff --git a/sdks/java/maven-archetypes/gcp-bom-examples/build.gradle b/sdks/java/maven-archetypes/gcp-bom-examples/build.gradle
index 541c91b..f9fabcf 100644
--- a/sdks/java/maven-archetypes/gcp-bom-examples/build.gradle
+++ b/sdks/java/maven-archetypes/gcp-bom-examples/build.gradle
@@ -71,6 +71,12 @@
     environment "HERE", "."
     commandLine '../examples/generate-sources.sh'
 }
+// add dependency BeamModulePlugin defined custom tasks
+// they are defined only when certain flags are provided (e.g. -Prelease; -Ppublishing, etc)
+def sourcesJar = project.tasks.findByName('sourcesJar')
+if (sourcesJar != null) {
+    sourcesJar.dependsOn generateSources
+}
 
 sourceSets {
     main {